From f707d88812cbbbdfbb596abff36628904a3403f3 Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Wed, 16 Sep 2026 17:07:04 +0200 Subject: [PATCH] Complete magnet holder navigation (#373) --- pyaml/common/holders/generic_array_holder.py | 51 ++++++++++ pyaml/common/holders/sub_holders.py | 50 ++++++++++ tests/common/test_array_holder_navigation.py | 99 ++++++++++++++++++++ tests/test_load_conf_with_code.py | 2 + 4 files changed, 202 insertions(+) create mode 100644 tests/common/test_array_holder_navigation.py diff --git a/pyaml/common/holders/generic_array_holder.py b/pyaml/common/holders/generic_array_holder.py index 0ce62441..4004a52a 100644 --- a/pyaml/common/holders/generic_array_holder.py +++ b/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) diff --git a/pyaml/common/holders/sub_holders.py b/pyaml/common/holders/sub_holders.py index 399485eb..d61088ba 100644 --- a/pyaml/common/holders/sub_holders.py +++ b/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"): diff --git a/tests/common/test_array_holder_navigation.py b/tests/common/test_array_holder_navigation.py new file mode 100644 index 00000000..3fdcbe42 --- /dev/null +++ b/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) diff --git a/tests/test_load_conf_with_code.py b/tests/test_load_conf_with_code.py index de476102..e65f5be5 100644 --- a/tests/test_load_conf_with_code.py +++ b/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