Complete magnet holder navigation (#373)

pyaml/common/holders/generic_array_holder.py

@@ -42,6 +42,13 @@ class GenericArrayHolder(Generic[T, A]):
         Return a named array or a transient array of all elements.
     add(arrayName, elementNames)
         Create and register a named array from element selectors.
+
+    Notes
+    -----
+    A configured array is also reachable as an attribute when its name is a
+    valid Python identifier, e.g. ``holder.QuadForTune`` is equivalent to
+    ``holder.get("QuadForTune")``. Array names appear in ``dir(holder)`` so
+    interactive completion (IPython, Jupyter) discovers them.
     """
 
     def __init__(
@@ -117,5 +124,49 @@ def __getitem__(self, key):
         """
         return self.get().__getitem__(key)
 
+    def __getattr__(self, name: str) -> A:
+        """
+        Return a configured array through attribute access.
+
+        Only called when normal attribute lookup fails, so it never shadows
+        :meth:`get`, :meth:`add`, or any other existing attribute.
+
+        Parameters
+        ----------
+        name : str
+            Configured array name. Must be a valid Python identifier.
+
+        Returns
+        -------
+        A
+            The array registered under ``name``.
+
+        Raises
+        ------
+        AttributeError
+            If ``name`` starts with an underscore, is not a valid Python
+            identifier, or does not match a configured array.
+
+        Examples
+        --------
+        >>> quad_family = sr.live.magnets.get("QuadForTune")
+        >>> same_quad_family = sr.live.magnets.QuadForTune
+        """
+        if name.startswith("_") or not name.isidentifier() or name not in self._array_store:
+            raise AttributeError(f"'{type(self).__name__}' object has no array named '{name}'")
+        return self._array_store[name]
+
+    def __dir__(self) -> list[str]:
+        """
+        List attributes, including configured array names.
+
+        Returns
+        -------
+        list of str
+            Default attributes plus configured array names that are valid
+            Python identifiers, for interactive completion (IPython, Jupyter).
+        """
+        return sorted(set(super().__dir__()) | {name for name in self._array_store if name.isidentifier()})
+
     def __repr__(self):
         return __pyaml_repr__(self)

pyaml/common/holders/sub_holders.py

@@ -42,6 +42,19 @@ class MagnetsHolder(GenericArrayHolder[Magnet, MagnetArray]):
     ----------
     peer : 'ElementHolder'
         Parent element holder that owns this specialized holder.
+
+    Notes
+    -----
+    A configured magnet family is also reachable as an attribute when its
+    name is a valid Python identifier, e.g. ``magnets.QuadForTune`` is
+    equivalent to ``magnets.get("QuadForTune")``. Family names appear in
+    ``dir(magnets)`` for interactive completion.
+
+    Examples
+    --------
+    >>> quad_family = sr.live.magnets.get("QuadForTune")
+    >>> same_quad_family = sr.live.magnets.QuadForTune
+    >>> combined_function_magnets = sr.live.magnets.get_cfm()
     """
 
     def __init__(self, peer: "ElementHolder"):
@@ -57,6 +70,22 @@ def __init__(self, peer: "ElementHolder"):
             "Magnet array",
         )
 
+    def get_cfm(self) -> CombinedFunctionMagnetArray:
+        """
+        Return all combined-function magnets.
+
+        Returns
+        -------
+        CombinedFunctionMagnetArray
+            New unnamed container with every registered combined-function
+            magnet.
+
+        Examples
+        --------
+        >>> combined_function_magnets = sr.live.magnets.get_cfm()
+        """
+        return self._peer.combined_function_magnets.get()
+
 
 class CombinedFunctionMagnetHolder(GenericElementHolder[CombinedFunctionMagnet]):
     """
@@ -83,6 +112,13 @@ class CombinedFunctionMagnetsHolder(GenericArrayHolder[CombinedFunctionMagnet, C
     ----------
     peer : 'ElementHolder'
         Parent element holder that owns this specialized holder.
+
+    Notes
+    -----
+    A configured array is also reachable as an attribute when its name is a
+    valid Python identifier, e.g. ``combined_function_magnets.CFM`` is
+    equivalent to ``combined_function_magnets.get("CFM")``. Array names
+    appear in ``dir(combined_function_magnets)`` for interactive completion.
     """
 
     def __init__(self, peer: "ElementHolder"):
@@ -124,6 +160,13 @@ class SerializedMagnetsHolder(GenericArrayHolder[SerializedMagnets, SerializedMa
     ----------
     peer : 'ElementHolder'
         Parent element holder that owns this specialized holder.
+
+    Notes
+    -----
+    A configured array is also reachable as an attribute when its name is a
+    valid Python identifier, e.g. ``serialized_magnets.QForTune`` is
+    equivalent to ``serialized_magnets.get("QForTune")``. Array names appear
+    in ``dir(serialized_magnets)`` for interactive completion.
     """
 
     def __init__(self, peer: "ElementHolder"):
@@ -165,6 +208,13 @@ class BPMsHolder(GenericArrayHolder[BPM, BPMArray]):
     ----------
     peer : 'ElementHolder'
         Parent element holder that owns this specialized holder.
+
+    Notes
+    -----
+    A configured array is also reachable as an attribute when its name is a
+    valid Python identifier, e.g. ``bpms.BPMS`` is equivalent to
+    ``bpms.get("BPMS")``. Array names appear in ``dir(bpms)`` for interactive
+    completion.
     """
 
     def __init__(self, peer: "ElementHolder"):

tests/common/test_array_holder_navigation.py

@@ -0,0 +1,99 @@
+import pytest
+
+from pyaml.accelerator import Accelerator
+from pyaml.arrays.cfm_magnet_array import CombinedFunctionMagnetArray
+
+
+@pytest.fixture
+def holder(accelerator_from_fragments, sr_configuration_fragments):
+    sr = accelerator_from_fragments(*sr_configuration_fragments)
+    sr.design.get_lattice().disable_6d()
+    return sr.design
+
+
+@pytest.fixture
+def serialized_holder():
+    sr = Accelerator.load("tests/config/sr_serialized_magnets.yaml", include_locations=False, ignore_external=True)
+    return sr.design
+
+
+def test_dynamic_attribute_returns_named_magnet_array(holder):
+    assert holder.magnets.HCORR is holder.magnets.get("HCORR")
+    assert holder.magnets.VCORR is holder.magnets.get("VCORR")
+    assert holder.magnets.HVCORR is holder.magnets.get("HVCORR")
+
+
+def test_dynamic_attribute_returns_named_combined_function_magnet_array(holder):
+    assert holder.combined_function_magnets.CFM is holder.combined_function_magnets.get("CFM")
+
+
+def test_dynamic_attribute_returns_named_bpm_array(holder):
+    assert holder.bpms.BPMS is holder.bpms.get("BPMS")
+
+
+def test_dynamic_attribute_returns_named_serialized_magnet_array(serialized_holder):
+    assert serialized_holder.serialized_magnets.QForTune is serialized_holder.serialized_magnets.get("QForTune")
+    assert serialized_holder.serialized_magnets.series is serialized_holder.serialized_magnets.get("series")
+
+
+def test_dynamic_attribute_unknown_name_raises_attribute_error(holder):
+    with pytest.raises(AttributeError):
+        _ = holder.magnets.UNKNOWN
+
+
+def test_dynamic_attribute_does_not_shadow_existing_methods(holder):
+    holder.magnets._array_store["get"] = holder.magnets.get("HCORR")
+    holder.magnets._array_store["add"] = holder.magnets.get("HCORR")
+
+    assert callable(holder.magnets.get)
+    assert callable(holder.magnets.add)
+    assert holder.magnets.get("HCORR").names() == ["SH1A-C01-H", "SH1A-C02-H"]
+
+
+def test_dynamic_attribute_rejects_non_identifier_names(holder):
+    holder.magnets._array_store["not-an-id"] = holder.magnets.get("HCORR")
+
+    with pytest.raises(AttributeError):
+        getattr(holder.magnets, "not-an-id")
+
+
+def test_dir_includes_configured_array_names(holder):
+    assert {"HCORR", "VCORR", "HVCORR"} <= set(dir(holder.magnets))
+    assert "CFM" in dir(holder.combined_function_magnets)
+    assert "BPMS" in dir(holder.bpms)
+
+
+def test_dir_excludes_non_identifier_array_names(holder):
+    holder.magnets._array_store["not-an-id"] = holder.magnets.get("HCORR")
+
+    assert "not-an-id" not in dir(holder.magnets)
+
+
+def test_get_cfm_returns_all_combined_function_magnets(holder):
+    combined_function_magnets = holder.magnets.get_cfm()
+
+    assert type(combined_function_magnets) is CombinedFunctionMagnetArray
+    assert combined_function_magnets.names() == holder.combined_function_magnets.get().names()
+
+
+def test_issue_373_example():
+    """Reproduce the #373 example verbatim, against real test lattice data."""
+    sr = Accelerator.load("tests/config/EBSTune-patterns.yaml", ignore_external=True)
+    sr.design.get_lattice().disable_6d()
+
+    all_magnets = sr.design.magnets[:]
+    quad_family = sr.design.magnets.get("QForTune")
+    same_quad_family = sr.design.magnets.QForTune
+    assert same_quad_family.names() == quad_family.names()
+    assert len(quad_family) == 124
+    assert len(all_magnets) >= len(quad_family)
+
+    combined_function_magnets = sr.design.magnets.get_cfm()
+    assert len(combined_function_magnets) == 0
+
+    matching_quadrupoles = sr.design.magnets["QF1*"]
+    assert matching_quadrupoles.names() == sr.design.magnets["QF1*"].names()
+
+    one_magnet = sr.design.magnet.get("QF1E-C04")
+    one_magnet.strength.set(0.8)
+    assert one_magnet.strength.get() == pytest.approx(0.8)

tests/test_load_conf_with_code.py

@@ -16,3 +16,5 @@ def test_load_conf_with_code():
     assert sr.live["BPM*"].names() == bpms.names()
     assert sr.live[:].names() == [element.get_name() for element in sr.live.get_all_elements()]
     assert sr.design["BPM*"].names() == sr.design.bpms.get("BPM").names()
+
+    assert sr.live.bpms.BPM is bpms