utils.py

pyat/at/lattice/utils.py

@@ -841,13 +841,7 @@ def get_value_refpts(ring: Sequence[Element], refpts: Refpts,
     Returns:
         attrvalues: numpy Array of attribute values.
     """
-    if index is None:
-        def getf(elem):
-            return getattr(elem, attrname)
-    else:
-        def getf(elem):
-            return getattr(elem, attrname)[index]
-
+    getf = getval(attrname, index=index)
     return numpy.array([getf(elem) for elem in refpts_iterator(ring, refpts,
                                                                regex=regex)])
 
@@ -884,13 +878,7 @@ def set_value_refpts(ring: Sequence[Element], refpts: Refpts,
              elements are shared with the original lattice.
              Any further modification will affect both lattices.
     """
-    if index is None:
-        def setf(elem, value):
-            setattr(elem, attrname, value)
-    else:
-        def setf(elem, value):
-            getattr(elem, attrname)[index] = value
-
+    setf = setval(attrname, index=index)
     if increment:
         attrvalues += get_value_refpts(ring, refpts,
                                        attrname, index=index,
@@ -903,8 +891,7 @@ def setf(elem, value):
     # noinspection PyShadowingNames
     @make_copy(copy)
     def apply(ring, refpts, values, regex):
-        for elm, val in zip(refpts_iterator(ring, refpts,
-                                            regex=regex), values):
+        for elm, val in zip(refpts_iterator(ring, refpts, regex=regex), values):
             setf(elm, val)
 
     return apply(ring, refpts, attrvalues, regex)