Add index to get_expr

xsuite · Oct 5, 2024 · 5675a5d · 5675a5d
1 parent fd64d87
commit 5675a5d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/xdeps/madxutils.py b/xdeps/madxutils.py
@@ -149,9 +149,15 @@ def get_value(self, key=None):
         else:
             return getattr(self._obj, key)
 
-    def get_expr(self, key=None):
+    def get_expr(self, key=None, index=None):
+        if index is not None:
+            if key is None:
+                raise ValueError('`key` must be provided when `index` is provided.')
+            return getattr(self._ref, key)[index]._expr
+
         if not hasattr(self._obj, "keys") and not hasattr(self._obj, "_xofields"):
-            raise ValueError("get_expr not supported for this object")
+            # is an array
+            return self._ref[key]._expr
 
         if key is None:
             if hasattr(self._obj, "_xofields"):