From 2ff348c26f9d6d6e10b86580b5611e9a5f73d1de Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 14:51:23 +0200 Subject: [PATCH 1/8] feat: add generic ElementHolder collection access - Return an element or None for exact-name lookups - Support typed wildcard selections, slices and insertion-order indexing - Add get() and [:] access to all registered elements - Reuse array type inference for selections - Document the API with NumPy docstrings and readable test examples --- pyaml/arrays/element_array.py | 35 ++++ pyaml/common/holders/element_holder.py | 91 ++++++++- tests/common/test_element_holder.py | 55 +++++ .../common/test_element_holder_collection.py | 192 ++++++++++++++++++ tests/test_load_conf_with_code.py | 5 + 5 files changed, 377 insertions(+), 1 deletion(-) create mode 100644 tests/common/test_element_holder.py create mode 100644 tests/common/test_element_holder_collection.py diff --git a/pyaml/arrays/element_array.py b/pyaml/arrays/element_array.py index 942ecd85..1d1f5006 100644 --- a/pyaml/arrays/element_array.py +++ b/pyaml/arrays/element_array.py @@ -130,6 +130,25 @@ def __auto_array(self, elements: list[Element]): if len(elements) == 0: return [] + return self._typed_array(elements) + + def _typed_array(self, elements: list[Element]) -> "ElementArray": + """Build a collection using the most specific compatible array type. + + Parameters + ---------- + elements : list[Element] + Selected references, in their desired order. + + Returns + ------- + ElementArray + Specialized array when possible, otherwise a generic array. + An empty selection returns an empty generic array. + """ + if not elements: + return self.__create_array("", Element, elements) + import inspect def mro_as_list(cls: type) -> list[type]: @@ -157,6 +176,22 @@ def mro_as_list(cls: type) -> list[type]: return self.__create_array("", chosen, elements) + def _select_names(self, pattern: str) -> "ElementArray": + """Select names without interpreting field selectors. + + Parameters + ---------- + pattern : str + A fnmatch pattern applied to each element name. + + Returns + ------- + ElementArray + Typed selection in the original order, including an empty array + when no names match. + """ + return self._typed_array([element for element in self if fnmatch.fnmatch(element.get_name(), pattern)]) + def __is_bool_mask(self, other: object) -> bool: """Return True if 'other' looks like a boolean mask (list or numpy array).""" # --- numpy boolean array --- diff --git a/pyaml/common/holders/element_holder.py b/pyaml/common/holders/element_holder.py index 0b472c9f..1b654053 100644 --- a/pyaml/common/holders/element_holder.py +++ b/pyaml/common/holders/element_holder.py @@ -5,7 +5,7 @@ import fnmatch import re from abc import ABCMeta, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload from ...arrays.element_array import ElementArray from ...bpm.bpm import BPM @@ -214,6 +214,95 @@ def _get(self, what, name, array) -> Element: return array[name] # Generic elements + def get(self) -> ElementArray: + """Return all registered elements in insertion order. + + Returns + ------- + ElementArray + New unnamed container sharing the registered element references. + + Notes + ----- + Registration order is not necessarily longitudinal lattice order. + Changing the returned container does not change the holder registry. + Each call reflects the current registry. + + Examples + -------- + >>> elements = sr.live.get() + >>> names = elements.names() + """ + return ElementArray("", list(self._ALL.values())) + + @overload + def __getitem__(self, key: int) -> Element: ... + + @overload + def __getitem__(self, key: slice) -> ElementArray: ... + + @overload + def __getitem__(self, key: str) -> Element | ElementArray | None: ... + + def __getitem__(self, key: int | slice | str) -> Element | ElementArray | None: + """Retrieve an element or select a collection. + + Parameters + ---------- + key : int, slice or str + Index in registration order, slice, exact name, or name pattern. + Strings containing ``*``, ``?`` or ``[`` use fnmatch matching. + Other strings are exact registry keys. Colons are literal. + + Returns + ------- + Element or ElementArray or None + An index returns an element. An exact name returns its element + or None. Patterns and slices return the most specific compatible + array, or an empty ElementArray when nothing matches. + The full slice ``[:]`` returns a generic ElementArray, like get(). + + Raises + ------ + IndexError + If the index is out of bounds. + TypeError + If the key is neither an integer, a slice, nor a string. + ValueError + If a slice has a zero step. + + Notes + ----- + Indices follow insertion order, not necessarily lattice order. + Collections share element references but do not modify the registry. + Field filters and regular expressions are not interpreted here. + + Examples + -------- + >>> bpm = sr.live["BPM01"] + >>> missing = sr.live["UNKNOWN"] # None + >>> bpms = sr.live["BPM*"] + >>> bpms = sr.live["BPM0[123]"] # BPM01, BPM02 or BPM03 + >>> bpms = sr.live["BPM0[1-3]"] # Same selection using a range + >>> quads = sr.live["Q[FD]*"] # Names starting with QF or QD + >>> bpms = sr.live["BPM0[!3]"] # One character after BPM0, except 3 + >>> first = sr.live[0] + >>> subset = sr.live[1:10] + >>> all_elements = sr.live[:] + """ + if isinstance(key, str): + if any(marker in key for marker in "*?["): + return self.get()._select_names(key) + return self._ALL.get(key) + if isinstance(key, int): + return list(self._ALL.values())[key] + if isinstance(key, slice): + elements = self.get() + if key == slice(None): + return elements + return elements._typed_array(list(elements)[key]) + raise TypeError("ElementHolder keys must be integers, slices or strings") + def fill_element_array(self, arrayName: str, elementNames: list[str]): self._fill_array( arrayName, diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py new file mode 100644 index 00000000..7a6548c2 --- /dev/null +++ b/tests/common/test_element_holder.py @@ -0,0 +1,55 @@ +from pyaml.common.element import Element +from pyaml.common.holders.element_holder import ElementHolder + + +class DispatchElement(Element): + def _fill_device(self, holder: ElementHolder) -> None: + holder.filled_elements.append(self) + + +class DispatchHolder(ElementHolder): + def __init__(self): + super().__init__() + self.filled_elements = [] + + def create_magnet_strength_aggregator(self, magnets): + return None + + def create_magnet_hardware_aggregator(self, magnets): + return None + + def create_bpm_aggregators(self, bpms): + return [None, None, None] + + def _fill_magnet(self, magnet) -> None: + raise AssertionError("Unexpected magnet dispatch") + + def _fill_combined_function_magnet(self, magnet) -> None: + raise AssertionError("Unexpected combined-function magnet dispatch") + + def _fill_serialized_magnets(self, magnets) -> None: + raise AssertionError("Unexpected serialized magnet dispatch") + + def _fill_bpm(self, bpm) -> None: + raise AssertionError("Unexpected BPM dispatch") + + def _fill_rf_plant(self, rf_plant) -> None: + raise AssertionError("Unexpected RF plant dispatch") + + def _fill_betatron_tune_monitor(self, monitor) -> None: + raise AssertionError("Unexpected betatron tune monitor dispatch") + + def _fill_tool(self, tool) -> None: + raise AssertionError("Unexpected tool dispatch") + + def _fill_unbound_element(self, element) -> None: + raise AssertionError("Unexpected unbound element dispatch") + + +def test_fill_device_delegates_to_elements(): + holder = DispatchHolder() + elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] + + holder._fill_device(elements) + + assert holder.filled_elements == elements diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py new file mode 100644 index 00000000..7495839f --- /dev/null +++ b/tests/common/test_element_holder_collection.py @@ -0,0 +1,192 @@ +import pytest + +from pyaml.arrays.bpm_array import BPMArray +from pyaml.arrays.element_array import ElementArray +from pyaml.arrays.magnet_array import MagnetArray +from pyaml.bpm.bpm import BPM +from pyaml.common.element import Element +from pyaml.common.exception import PyAMLException +from pyaml.common.holders.element_holder import ElementHolder +from pyaml.magnet.magnet import Magnet + + +class CollectionHolder(ElementHolder): + """Minimal holder for collection examples, without device access.""" + + def __init__(self, *elements): + super().__init__() + for element in elements: + element._peer = self + self.add_element(element) + + def create_magnet_strength_aggregator(self, magnets): + return None + + def create_magnet_hardware_aggregator(self, magnets): + return None + + def create_bpm_aggregators(self, bpms): + return [None, None, None] + + +class Quadrupole(Magnet): + pass + + +class Corrector(Magnet): + pass + + +def test_exact_name_returns_the_element_or_none(): + bpm = BPM("BPM01") + holder = CollectionHolder(bpm) + + assert holder["BPM01"] is bpm + assert holder["UNKNOWN"] is None + + +def test_get_and_full_slice_keep_registration_order(): + holder = CollectionHolder(BPM("BPM02"), BPM("BPM01")) + + assert type(holder.get()) is ElementArray + assert type(holder[:]) is ElementArray + assert holder.get().names() == ["BPM02", "BPM01"] + assert holder[:].names() == ["BPM02", "BPM01"] + + +def test_collections_are_independent_but_share_elements(): + bpm = BPM("BPM01") + holder = CollectionHolder(bpm) + collection = holder.get() + + assert collection[0] is bpm + collection.clear() + assert holder["BPM01"] is bpm + assert holder.get().names() == ["BPM01"] + + +def test_new_calls_reflect_registry_additions(): + holder = CollectionHolder(BPM("BPM01")) + previous = holder.get() + added = BPM("BPM02").attach(holder, None, None, None) + holder.add_element(added) + + assert previous.names() == ["BPM01"] + assert holder.get().names() == ["BPM01", "BPM02"] + + +def test_patterns_always_return_arrays(): + holder = CollectionHolder(BPM("BPM02"), BPM("BPM01"), Element("MARKER")) + + assert isinstance(holder["BPM*"], BPMArray) + assert holder["BPM*"].names() == ["BPM02", "BPM01"] + assert isinstance(holder["BPM0[1]"], BPMArray) + assert holder["BPM0[1]"].names() == ["BPM01"] + assert holder["BPM0?"].names() == ["BPM02", "BPM01"] + assert type(holder["MISSING*"]) is ElementArray + assert holder["MISSING*"].names() == [] + + +def test_character_classes_in_name_patterns(): + holder = CollectionHolder( + BPM("BPM01"), + BPM("BPM02"), + BPM("BPM03"), + BPM("BPM04"), + Magnet("QF01"), + Magnet("QD01"), + Magnet("QS01"), + ) + + assert holder["BPM0[123]"].names() == ["BPM01", "BPM02", "BPM03"] + assert holder["BPM0[1-3]"].names() == ["BPM01", "BPM02", "BPM03"] + assert holder["Q[FD]*"].names() == ["QF01", "QD01"] + assert holder["BPM0[!3]"].names() == ["BPM01", "BPM02", "BPM04"] + + +def test_colons_are_part_of_names(): + bpm = BPM("CELL01:BPM01") + holder = CollectionHolder(bpm, Magnet("MAGNET")) + + assert holder["CELL01:BPM01"] is bpm + assert holder["CELL01:BPM*"].names() == ["CELL01:BPM01"] + assert holder["model_name:*"].names() == [] + + +def test_magnet_subclasses_share_a_typed_array_in_either_order(): + holder = CollectionHolder(Quadrupole("MAG_Q"), Corrector("MAG_C")) + + assert isinstance(holder["MAG*"], MagnetArray) + assert isinstance(holder[::-1], MagnetArray) + assert holder[::-1].names() == ["MAG_C", "MAG_Q"] + + +def test_mixed_selection_returns_a_generic_array(): + holder = CollectionHolder(Magnet("CELL_Q"), BPM("CELL_BPM")) + + assert type(holder["CELL*"]) is ElementArray + assert holder["CELL*"].names() == ["CELL_Q", "CELL_BPM"] + + +def test_indices_and_slices_follow_insertion_order(): + second = BPM("BPM02") + first = BPM("BPM01") + third = BPM("BPM03") + holder = CollectionHolder(second, first, third) + + assert holder[0] is second + assert holder[-1] is third + assert isinstance(holder[1:3], BPMArray) + assert holder[1:3].names() == ["BPM01", "BPM03"] + assert holder[::2].names() == ["BPM02", "BPM03"] + assert type(holder[3:]) is ElementArray + + +def test_empty_holder_returns_empty_collections(): + holder = CollectionHolder() + + assert type(holder.get()) is ElementArray + assert holder[:].names() == [] + assert holder["BPM*"].names() == [] + assert holder["BPM01"] is None + + +def test_invalid_indices_and_keys_raise_clear_errors(): + holder = CollectionHolder(BPM("BPM01")) + + with pytest.raises(IndexError): + holder[1] + with pytest.raises(IndexError): + holder[-2] + with pytest.raises(TypeError): + holder[1.5] + with pytest.raises(ValueError): + holder[::0] + + +def test_selection_intersects_with_a_configured_family(): + holder = CollectionHolder(Magnet("Q02"), Magnet("Q01"), BPM("BPM01")) + holder.fill_element_array("CELL01", ["Q01", "BPM01"]) + + selected = holder["Q*"] & holder.get_elements("CELL01") + + assert isinstance(selected, MagnetArray) + assert selected.names() == ["Q01"] + assert holder.get_element("Q01") is holder["Q01"] + assert holder.get_all_elements() == list(holder.get()) + with pytest.raises(PyAMLException): + holder.get_element("UNKNOWN") + + +def test_existing_array_field_filters_still_work(): + magnet = Magnet("VIRTUAL_Q") + magnet.set_model_name("PHYSICAL_Q") + holder = CollectionHolder(magnet) + + assert holder["VIRTUAL*"]["model_name:PHYSICAL*"].names() == ["VIRTUAL_Q"] + + +def test_existing_empty_intersection_stays_a_list(): + holder = CollectionHolder(Magnet("Q01"), Magnet("Q02")) + + assert type(holder["Q0[1]"] & holder["Q0[2]"]) is list diff --git a/tests/test_load_conf_with_code.py b/tests/test_load_conf_with_code.py index 42286ba2..de476102 100644 --- a/tests/test_load_conf_with_code.py +++ b/tests/test_load_conf_with_code.py @@ -11,3 +11,8 @@ def test_load_conf_with_code(): bpms = sr.live.bpms.get("BPM") assert bpms is not None assert len(bpms) == 320 + + assert sr.live[bpms[0].get_name()] is bpms[0] + 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() From 8682539e66a476c861219f982f3762d52eb6d25f Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 16:02:57 +0200 Subject: [PATCH 2/8] Test correction --- tests/common/test_element_holder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py index 7a6548c2..d7f66540 100644 --- a/tests/common/test_element_holder.py +++ b/tests/common/test_element_holder.py @@ -50,6 +50,6 @@ def test_fill_device_delegates_to_elements(): holder = DispatchHolder() elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] - holder._fill_device(elements) + holder.fill_device(elements) assert holder.filled_elements == elements From 211ee32a76d48c6ab3cc38d0afcb3c2d76434aed Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 16:12:10 +0200 Subject: [PATCH 3/8] Test correction --- tests/common/test_element_holder.py | 55 ----------------------------- 1 file changed, 55 deletions(-) delete mode 100644 tests/common/test_element_holder.py diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py deleted file mode 100644 index d7f66540..00000000 --- a/tests/common/test_element_holder.py +++ /dev/null @@ -1,55 +0,0 @@ -from pyaml.common.element import Element -from pyaml.common.holders.element_holder import ElementHolder - - -class DispatchElement(Element): - def _fill_device(self, holder: ElementHolder) -> None: - holder.filled_elements.append(self) - - -class DispatchHolder(ElementHolder): - def __init__(self): - super().__init__() - self.filled_elements = [] - - def create_magnet_strength_aggregator(self, magnets): - return None - - def create_magnet_hardware_aggregator(self, magnets): - return None - - def create_bpm_aggregators(self, bpms): - return [None, None, None] - - def _fill_magnet(self, magnet) -> None: - raise AssertionError("Unexpected magnet dispatch") - - def _fill_combined_function_magnet(self, magnet) -> None: - raise AssertionError("Unexpected combined-function magnet dispatch") - - def _fill_serialized_magnets(self, magnets) -> None: - raise AssertionError("Unexpected serialized magnet dispatch") - - def _fill_bpm(self, bpm) -> None: - raise AssertionError("Unexpected BPM dispatch") - - def _fill_rf_plant(self, rf_plant) -> None: - raise AssertionError("Unexpected RF plant dispatch") - - def _fill_betatron_tune_monitor(self, monitor) -> None: - raise AssertionError("Unexpected betatron tune monitor dispatch") - - def _fill_tool(self, tool) -> None: - raise AssertionError("Unexpected tool dispatch") - - def _fill_unbound_element(self, element) -> None: - raise AssertionError("Unexpected unbound element dispatch") - - -def test_fill_device_delegates_to_elements(): - holder = DispatchHolder() - elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] - - holder.fill_device(elements) - - assert holder.filled_elements == elements From 518b3dee23237667e97ffb0b0cfc8d8380296e49 Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Tue, 8 Sep 2026 10:51:53 +0200 Subject: [PATCH 4/8] Test improvements to use existing lattices --- .../common/test_element_holder_collection.py | 196 +++++++----------- 1 file changed, 77 insertions(+), 119 deletions(-) diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py index 7495839f..937d2da7 100644 --- a/tests/common/test_element_holder_collection.py +++ b/tests/common/test_element_holder_collection.py @@ -4,189 +4,147 @@ from pyaml.arrays.element_array import ElementArray from pyaml.arrays.magnet_array import MagnetArray from pyaml.bpm.bpm import BPM -from pyaml.common.element import Element from pyaml.common.exception import PyAMLException -from pyaml.common.holders.element_holder import ElementHolder -from pyaml.magnet.magnet import Magnet +from pyaml.lattice.simulator import Simulator -class CollectionHolder(ElementHolder): - """Minimal holder for collection examples, without device access.""" +@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 - def __init__(self, *elements): - super().__init__() - for element in elements: - element._peer = self - self.add_element(element) - def create_magnet_strength_aggregator(self, magnets): - return None - - def create_magnet_hardware_aggregator(self, magnets): - return None - - def create_bpm_aggregators(self, bpms): - return [None, None, None] - - -class Quadrupole(Magnet): - pass - - -class Corrector(Magnet): - pass - - -def test_exact_name_returns_the_element_or_none(): - bpm = BPM("BPM01") - holder = CollectionHolder(bpm) - - assert holder["BPM01"] is bpm +def test_exact_name_returns_the_element_or_none(holder): + assert holder["BPM_C04-01"] is holder.bpm.get("BPM_C04-01") assert holder["UNKNOWN"] is None -def test_get_and_full_slice_keep_registration_order(): - holder = CollectionHolder(BPM("BPM02"), BPM("BPM01")) +def test_get_and_full_slice_keep_registration_order(holder): + names = [element.get_name() for element in holder.get_all_elements()] assert type(holder.get()) is ElementArray assert type(holder[:]) is ElementArray - assert holder.get().names() == ["BPM02", "BPM01"] - assert holder[:].names() == ["BPM02", "BPM01"] + assert holder.get().names() == names + assert holder[:].names() == names + assert names != sorted(names) -def test_collections_are_independent_but_share_elements(): - bpm = BPM("BPM01") - holder = CollectionHolder(bpm) +def test_collections_are_independent_but_share_elements(holder): collection = holder.get() + first = holder[0] + original_names = collection.names() - assert collection[0] is bpm + assert collection[0] is first collection.clear() - assert holder["BPM01"] is bpm - assert holder.get().names() == ["BPM01"] + assert holder[0] is first + assert holder.get().names() == original_names -def test_new_calls_reflect_registry_additions(): - holder = CollectionHolder(BPM("BPM01")) +def test_new_calls_reflect_registry_additions(holder): previous = holder.get() - added = BPM("BPM02").attach(holder, None, None, None) - holder.add_element(added) - - assert previous.names() == ["BPM01"] - assert holder.get().names() == ["BPM01", "BPM02"] + holder.fill_device([BPM("EXTRA_BPM", lattice_names="list(BPM_C04-01)")]) + assert "EXTRA_BPM" not in previous.names() + assert holder.get().names() == previous.names() + ["EXTRA_BPM"] + assert holder["EXTRA_BPM"] is holder.bpm.get("EXTRA_BPM") -def test_patterns_always_return_arrays(): - holder = CollectionHolder(BPM("BPM02"), BPM("BPM01"), Element("MARKER")) +def test_patterns_always_return_arrays(holder): assert isinstance(holder["BPM*"], BPMArray) - assert holder["BPM*"].names() == ["BPM02", "BPM01"] - assert isinstance(holder["BPM0[1]"], BPMArray) - assert holder["BPM0[1]"].names() == ["BPM01"] - assert holder["BPM0?"].names() == ["BPM02", "BPM01"] + assert holder["BPM*"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert isinstance(holder["BPM_C04-0[1]"], BPMArray) + assert holder["BPM_C04-0[1]"].names() == ["BPM_C04-01"] + assert holder["BPM_C04-0?"].names() == ["BPM_C04-01", "BPM_C04-02"] assert type(holder["MISSING*"]) is ElementArray assert holder["MISSING*"].names() == [] -def test_character_classes_in_name_patterns(): - holder = CollectionHolder( - BPM("BPM01"), - BPM("BPM02"), - BPM("BPM03"), - BPM("BPM04"), - Magnet("QF01"), - Magnet("QD01"), - Magnet("QS01"), - ) +def test_character_classes_in_name_patterns(holder): + assert holder["BPM_C04-0[12]"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert holder["BPM_C04-0[1-2]"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert holder["SH1A-C01-[HV]*"].names() == ["SH1A-C01-H", "SH1A-C01-V"] + assert holder["BPM_C04-0[!2]"].names() == ["BPM_C04-01"] - assert holder["BPM0[123]"].names() == ["BPM01", "BPM02", "BPM03"] - assert holder["BPM0[1-3]"].names() == ["BPM01", "BPM02", "BPM03"] - assert holder["Q[FD]*"].names() == ["QF01", "QD01"] - assert holder["BPM0[!3]"].names() == ["BPM01", "BPM02", "BPM04"] +def test_colons_are_part_of_names(holder): + holder.fill_device([BPM("CELL04:BPM01", lattice_names="list(BPM_C04-01)")]) -def test_colons_are_part_of_names(): - bpm = BPM("CELL01:BPM01") - holder = CollectionHolder(bpm, Magnet("MAGNET")) - - assert holder["CELL01:BPM01"] is bpm - assert holder["CELL01:BPM*"].names() == ["CELL01:BPM01"] + assert holder["CELL04:BPM01"] is holder.bpm.get("CELL04:BPM01") + assert holder["CELL04:BPM*"].names() == ["CELL04:BPM01"] assert holder["model_name:*"].names() == [] -def test_magnet_subclasses_share_a_typed_array_in_either_order(): - holder = CollectionHolder(Quadrupole("MAG_Q"), Corrector("MAG_C")) +def test_magnet_subclasses_share_a_typed_array_in_either_order(holder): + horizontal = holder.magnet.get("SH1A-C01-H") + vertical = holder.magnet.get("SH1A-C01-V") + start = holder.get_all_elements().index(horizontal) - assert isinstance(holder["MAG*"], MagnetArray) - assert isinstance(holder[::-1], MagnetArray) - assert holder[::-1].names() == ["MAG_C", "MAG_Q"] + assert type(horizontal) is not type(vertical) + assert isinstance(holder["SH1A-C01-[HV]"], MagnetArray) + assert isinstance(holder[start : start + 2], MagnetArray) + assert isinstance(holder[start + 1 : start - 1 : -1], MagnetArray) + assert holder[start + 1 : start - 1 : -1].names() == ["SH1A-C01-V", "SH1A-C01-H"] -def test_mixed_selection_returns_a_generic_array(): - holder = CollectionHolder(Magnet("CELL_Q"), BPM("CELL_BPM")) +def test_mixed_selection_returns_a_generic_array(holder): + selected = holder["*-C01*"] - assert type(holder["CELL*"]) is ElementArray - assert holder["CELL*"].names() == ["CELL_Q", "CELL_BPM"] + assert type(selected) is ElementArray + assert "QF1A-C01" in selected.names() + assert "SH1A-C01" in selected.names() -def test_indices_and_slices_follow_insertion_order(): - second = BPM("BPM02") - first = BPM("BPM01") - third = BPM("BPM03") - holder = CollectionHolder(second, first, third) +def test_indices_and_slices_follow_insertion_order(holder): + registered = holder.get_all_elements() - assert holder[0] is second - assert holder[-1] is third - assert isinstance(holder[1:3], BPMArray) - assert holder[1:3].names() == ["BPM01", "BPM03"] - assert holder[::2].names() == ["BPM02", "BPM03"] - assert type(holder[3:]) is ElementArray + assert holder[0] is registered[0] + assert holder[-1] is registered[-1] + assert list(holder[1:3]) == registered[1:3] + assert list(holder[::2]) == registered[::2] + assert isinstance(holder[-2:], BPMArray) + assert holder[-2:].names() == ["BPM_C04-01", "BPM_C04-02"] + assert type(holder[len(registered) :]) is ElementArray -def test_empty_holder_returns_empty_collections(): - holder = CollectionHolder() +def test_empty_holder_returns_empty_collections(ebs_lattice_file): + holder = Simulator(name="empty", lattice=str(ebs_lattice_file)) assert type(holder.get()) is ElementArray assert holder[:].names() == [] assert holder["BPM*"].names() == [] - assert holder["BPM01"] is None + assert holder["BPM_C04-01"] is None -def test_invalid_indices_and_keys_raise_clear_errors(): - holder = CollectionHolder(BPM("BPM01")) +def test_invalid_indices_and_keys_raise_clear_errors(holder): + size = len(holder.get_all_elements()) with pytest.raises(IndexError): - holder[1] + holder[size] with pytest.raises(IndexError): - holder[-2] + holder[-size - 1] with pytest.raises(TypeError): holder[1.5] with pytest.raises(ValueError): holder[::0] -def test_selection_intersects_with_a_configured_family(): - holder = CollectionHolder(Magnet("Q02"), Magnet("Q01"), BPM("BPM01")) - holder.fill_element_array("CELL01", ["Q01", "BPM01"]) - - selected = holder["Q*"] & holder.get_elements("CELL01") +def test_selection_intersects_with_a_configured_family(holder): + selected = holder["SH1A-C0?-H"] & holder.get_elements("ElArray") assert isinstance(selected, MagnetArray) - assert selected.names() == ["Q01"] - assert holder.get_element("Q01") is holder["Q01"] + assert selected.names() == ["SH1A-C02-H"] + assert holder.get_element("SH1A-C02-H") is holder["SH1A-C02-H"] assert holder.get_all_elements() == list(holder.get()) with pytest.raises(PyAMLException): holder.get_element("UNKNOWN") -def test_existing_array_field_filters_still_work(): - magnet = Magnet("VIRTUAL_Q") - magnet.set_model_name("PHYSICAL_Q") - holder = CollectionHolder(magnet) - - assert holder["VIRTUAL*"]["model_name:PHYSICAL*"].names() == ["VIRTUAL_Q"] +def test_existing_array_field_filters_still_work(holder): + selected = holder["SH1A-C0?-H"]["model_name:SH1A-C01"] + assert selected.names() == ["SH1A-C01-H"] -def test_existing_empty_intersection_stays_a_list(): - holder = CollectionHolder(Magnet("Q01"), Magnet("Q02")) - assert type(holder["Q0[1]"] & holder["Q0[2]"]) is list +def test_existing_empty_intersection_stays_a_list(holder): + assert type(holder["SH1A-C0?-H"] & holder["SH1A-C0?-V"]) is list From 5b0a5c079f7899099bc56ba1d24c480284b2da3f Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 14:51:23 +0200 Subject: [PATCH 5/8] feat: add generic ElementHolder collection access - Return an element or None for exact-name lookups - Support typed wildcard selections, slices and insertion-order indexing - Add get() and [:] access to all registered elements - Reuse array type inference for selections - Document the API with NumPy docstrings and readable test examples --- pyaml/arrays/element_array.py | 35 ++++ pyaml/common/holders/element_holder.py | 91 ++++++++- tests/common/test_element_holder.py | 55 +++++ .../common/test_element_holder_collection.py | 192 ++++++++++++++++++ tests/test_load_conf_with_code.py | 5 + 5 files changed, 377 insertions(+), 1 deletion(-) create mode 100644 tests/common/test_element_holder.py create mode 100644 tests/common/test_element_holder_collection.py diff --git a/pyaml/arrays/element_array.py b/pyaml/arrays/element_array.py index 942ecd85..1d1f5006 100644 --- a/pyaml/arrays/element_array.py +++ b/pyaml/arrays/element_array.py @@ -130,6 +130,25 @@ def __auto_array(self, elements: list[Element]): if len(elements) == 0: return [] + return self._typed_array(elements) + + def _typed_array(self, elements: list[Element]) -> "ElementArray": + """Build a collection using the most specific compatible array type. + + Parameters + ---------- + elements : list[Element] + Selected references, in their desired order. + + Returns + ------- + ElementArray + Specialized array when possible, otherwise a generic array. + An empty selection returns an empty generic array. + """ + if not elements: + return self.__create_array("", Element, elements) + import inspect def mro_as_list(cls: type) -> list[type]: @@ -157,6 +176,22 @@ def mro_as_list(cls: type) -> list[type]: return self.__create_array("", chosen, elements) + def _select_names(self, pattern: str) -> "ElementArray": + """Select names without interpreting field selectors. + + Parameters + ---------- + pattern : str + A fnmatch pattern applied to each element name. + + Returns + ------- + ElementArray + Typed selection in the original order, including an empty array + when no names match. + """ + return self._typed_array([element for element in self if fnmatch.fnmatch(element.get_name(), pattern)]) + def __is_bool_mask(self, other: object) -> bool: """Return True if 'other' looks like a boolean mask (list or numpy array).""" # --- numpy boolean array --- diff --git a/pyaml/common/holders/element_holder.py b/pyaml/common/holders/element_holder.py index 0b472c9f..1b654053 100644 --- a/pyaml/common/holders/element_holder.py +++ b/pyaml/common/holders/element_holder.py @@ -5,7 +5,7 @@ import fnmatch import re from abc import ABCMeta, abstractmethod -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload from ...arrays.element_array import ElementArray from ...bpm.bpm import BPM @@ -214,6 +214,95 @@ def _get(self, what, name, array) -> Element: return array[name] # Generic elements + def get(self) -> ElementArray: + """Return all registered elements in insertion order. + + Returns + ------- + ElementArray + New unnamed container sharing the registered element references. + + Notes + ----- + Registration order is not necessarily longitudinal lattice order. + Changing the returned container does not change the holder registry. + Each call reflects the current registry. + + Examples + -------- + >>> elements = sr.live.get() + >>> names = elements.names() + """ + return ElementArray("", list(self._ALL.values())) + + @overload + def __getitem__(self, key: int) -> Element: ... + + @overload + def __getitem__(self, key: slice) -> ElementArray: ... + + @overload + def __getitem__(self, key: str) -> Element | ElementArray | None: ... + + def __getitem__(self, key: int | slice | str) -> Element | ElementArray | None: + """Retrieve an element or select a collection. + + Parameters + ---------- + key : int, slice or str + Index in registration order, slice, exact name, or name pattern. + Strings containing ``*``, ``?`` or ``[`` use fnmatch matching. + Other strings are exact registry keys. Colons are literal. + + Returns + ------- + Element or ElementArray or None + An index returns an element. An exact name returns its element + or None. Patterns and slices return the most specific compatible + array, or an empty ElementArray when nothing matches. + The full slice ``[:]`` returns a generic ElementArray, like get(). + + Raises + ------ + IndexError + If the index is out of bounds. + TypeError + If the key is neither an integer, a slice, nor a string. + ValueError + If a slice has a zero step. + + Notes + ----- + Indices follow insertion order, not necessarily lattice order. + Collections share element references but do not modify the registry. + Field filters and regular expressions are not interpreted here. + + Examples + -------- + >>> bpm = sr.live["BPM01"] + >>> missing = sr.live["UNKNOWN"] # None + >>> bpms = sr.live["BPM*"] + >>> bpms = sr.live["BPM0[123]"] # BPM01, BPM02 or BPM03 + >>> bpms = sr.live["BPM0[1-3]"] # Same selection using a range + >>> quads = sr.live["Q[FD]*"] # Names starting with QF or QD + >>> bpms = sr.live["BPM0[!3]"] # One character after BPM0, except 3 + >>> first = sr.live[0] + >>> subset = sr.live[1:10] + >>> all_elements = sr.live[:] + """ + if isinstance(key, str): + if any(marker in key for marker in "*?["): + return self.get()._select_names(key) + return self._ALL.get(key) + if isinstance(key, int): + return list(self._ALL.values())[key] + if isinstance(key, slice): + elements = self.get() + if key == slice(None): + return elements + return elements._typed_array(list(elements)[key]) + raise TypeError("ElementHolder keys must be integers, slices or strings") + def fill_element_array(self, arrayName: str, elementNames: list[str]): self._fill_array( arrayName, diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py new file mode 100644 index 00000000..7a6548c2 --- /dev/null +++ b/tests/common/test_element_holder.py @@ -0,0 +1,55 @@ +from pyaml.common.element import Element +from pyaml.common.holders.element_holder import ElementHolder + + +class DispatchElement(Element): + def _fill_device(self, holder: ElementHolder) -> None: + holder.filled_elements.append(self) + + +class DispatchHolder(ElementHolder): + def __init__(self): + super().__init__() + self.filled_elements = [] + + def create_magnet_strength_aggregator(self, magnets): + return None + + def create_magnet_hardware_aggregator(self, magnets): + return None + + def create_bpm_aggregators(self, bpms): + return [None, None, None] + + def _fill_magnet(self, magnet) -> None: + raise AssertionError("Unexpected magnet dispatch") + + def _fill_combined_function_magnet(self, magnet) -> None: + raise AssertionError("Unexpected combined-function magnet dispatch") + + def _fill_serialized_magnets(self, magnets) -> None: + raise AssertionError("Unexpected serialized magnet dispatch") + + def _fill_bpm(self, bpm) -> None: + raise AssertionError("Unexpected BPM dispatch") + + def _fill_rf_plant(self, rf_plant) -> None: + raise AssertionError("Unexpected RF plant dispatch") + + def _fill_betatron_tune_monitor(self, monitor) -> None: + raise AssertionError("Unexpected betatron tune monitor dispatch") + + def _fill_tool(self, tool) -> None: + raise AssertionError("Unexpected tool dispatch") + + def _fill_unbound_element(self, element) -> None: + raise AssertionError("Unexpected unbound element dispatch") + + +def test_fill_device_delegates_to_elements(): + holder = DispatchHolder() + elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] + + holder._fill_device(elements) + + assert holder.filled_elements == elements diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py new file mode 100644 index 00000000..7495839f --- /dev/null +++ b/tests/common/test_element_holder_collection.py @@ -0,0 +1,192 @@ +import pytest + +from pyaml.arrays.bpm_array import BPMArray +from pyaml.arrays.element_array import ElementArray +from pyaml.arrays.magnet_array import MagnetArray +from pyaml.bpm.bpm import BPM +from pyaml.common.element import Element +from pyaml.common.exception import PyAMLException +from pyaml.common.holders.element_holder import ElementHolder +from pyaml.magnet.magnet import Magnet + + +class CollectionHolder(ElementHolder): + """Minimal holder for collection examples, without device access.""" + + def __init__(self, *elements): + super().__init__() + for element in elements: + element._peer = self + self.add_element(element) + + def create_magnet_strength_aggregator(self, magnets): + return None + + def create_magnet_hardware_aggregator(self, magnets): + return None + + def create_bpm_aggregators(self, bpms): + return [None, None, None] + + +class Quadrupole(Magnet): + pass + + +class Corrector(Magnet): + pass + + +def test_exact_name_returns_the_element_or_none(): + bpm = BPM("BPM01") + holder = CollectionHolder(bpm) + + assert holder["BPM01"] is bpm + assert holder["UNKNOWN"] is None + + +def test_get_and_full_slice_keep_registration_order(): + holder = CollectionHolder(BPM("BPM02"), BPM("BPM01")) + + assert type(holder.get()) is ElementArray + assert type(holder[:]) is ElementArray + assert holder.get().names() == ["BPM02", "BPM01"] + assert holder[:].names() == ["BPM02", "BPM01"] + + +def test_collections_are_independent_but_share_elements(): + bpm = BPM("BPM01") + holder = CollectionHolder(bpm) + collection = holder.get() + + assert collection[0] is bpm + collection.clear() + assert holder["BPM01"] is bpm + assert holder.get().names() == ["BPM01"] + + +def test_new_calls_reflect_registry_additions(): + holder = CollectionHolder(BPM("BPM01")) + previous = holder.get() + added = BPM("BPM02").attach(holder, None, None, None) + holder.add_element(added) + + assert previous.names() == ["BPM01"] + assert holder.get().names() == ["BPM01", "BPM02"] + + +def test_patterns_always_return_arrays(): + holder = CollectionHolder(BPM("BPM02"), BPM("BPM01"), Element("MARKER")) + + assert isinstance(holder["BPM*"], BPMArray) + assert holder["BPM*"].names() == ["BPM02", "BPM01"] + assert isinstance(holder["BPM0[1]"], BPMArray) + assert holder["BPM0[1]"].names() == ["BPM01"] + assert holder["BPM0?"].names() == ["BPM02", "BPM01"] + assert type(holder["MISSING*"]) is ElementArray + assert holder["MISSING*"].names() == [] + + +def test_character_classes_in_name_patterns(): + holder = CollectionHolder( + BPM("BPM01"), + BPM("BPM02"), + BPM("BPM03"), + BPM("BPM04"), + Magnet("QF01"), + Magnet("QD01"), + Magnet("QS01"), + ) + + assert holder["BPM0[123]"].names() == ["BPM01", "BPM02", "BPM03"] + assert holder["BPM0[1-3]"].names() == ["BPM01", "BPM02", "BPM03"] + assert holder["Q[FD]*"].names() == ["QF01", "QD01"] + assert holder["BPM0[!3]"].names() == ["BPM01", "BPM02", "BPM04"] + + +def test_colons_are_part_of_names(): + bpm = BPM("CELL01:BPM01") + holder = CollectionHolder(bpm, Magnet("MAGNET")) + + assert holder["CELL01:BPM01"] is bpm + assert holder["CELL01:BPM*"].names() == ["CELL01:BPM01"] + assert holder["model_name:*"].names() == [] + + +def test_magnet_subclasses_share_a_typed_array_in_either_order(): + holder = CollectionHolder(Quadrupole("MAG_Q"), Corrector("MAG_C")) + + assert isinstance(holder["MAG*"], MagnetArray) + assert isinstance(holder[::-1], MagnetArray) + assert holder[::-1].names() == ["MAG_C", "MAG_Q"] + + +def test_mixed_selection_returns_a_generic_array(): + holder = CollectionHolder(Magnet("CELL_Q"), BPM("CELL_BPM")) + + assert type(holder["CELL*"]) is ElementArray + assert holder["CELL*"].names() == ["CELL_Q", "CELL_BPM"] + + +def test_indices_and_slices_follow_insertion_order(): + second = BPM("BPM02") + first = BPM("BPM01") + third = BPM("BPM03") + holder = CollectionHolder(second, first, third) + + assert holder[0] is second + assert holder[-1] is third + assert isinstance(holder[1:3], BPMArray) + assert holder[1:3].names() == ["BPM01", "BPM03"] + assert holder[::2].names() == ["BPM02", "BPM03"] + assert type(holder[3:]) is ElementArray + + +def test_empty_holder_returns_empty_collections(): + holder = CollectionHolder() + + assert type(holder.get()) is ElementArray + assert holder[:].names() == [] + assert holder["BPM*"].names() == [] + assert holder["BPM01"] is None + + +def test_invalid_indices_and_keys_raise_clear_errors(): + holder = CollectionHolder(BPM("BPM01")) + + with pytest.raises(IndexError): + holder[1] + with pytest.raises(IndexError): + holder[-2] + with pytest.raises(TypeError): + holder[1.5] + with pytest.raises(ValueError): + holder[::0] + + +def test_selection_intersects_with_a_configured_family(): + holder = CollectionHolder(Magnet("Q02"), Magnet("Q01"), BPM("BPM01")) + holder.fill_element_array("CELL01", ["Q01", "BPM01"]) + + selected = holder["Q*"] & holder.get_elements("CELL01") + + assert isinstance(selected, MagnetArray) + assert selected.names() == ["Q01"] + assert holder.get_element("Q01") is holder["Q01"] + assert holder.get_all_elements() == list(holder.get()) + with pytest.raises(PyAMLException): + holder.get_element("UNKNOWN") + + +def test_existing_array_field_filters_still_work(): + magnet = Magnet("VIRTUAL_Q") + magnet.set_model_name("PHYSICAL_Q") + holder = CollectionHolder(magnet) + + assert holder["VIRTUAL*"]["model_name:PHYSICAL*"].names() == ["VIRTUAL_Q"] + + +def test_existing_empty_intersection_stays_a_list(): + holder = CollectionHolder(Magnet("Q01"), Magnet("Q02")) + + assert type(holder["Q0[1]"] & holder["Q0[2]"]) is list diff --git a/tests/test_load_conf_with_code.py b/tests/test_load_conf_with_code.py index 42286ba2..de476102 100644 --- a/tests/test_load_conf_with_code.py +++ b/tests/test_load_conf_with_code.py @@ -11,3 +11,8 @@ def test_load_conf_with_code(): bpms = sr.live.bpms.get("BPM") assert bpms is not None assert len(bpms) == 320 + + assert sr.live[bpms[0].get_name()] is bpms[0] + 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() From afd1146a5132d2887a768f110888c98a2e9e0377 Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 16:02:57 +0200 Subject: [PATCH 6/8] Test correction --- tests/common/test_element_holder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py index 7a6548c2..d7f66540 100644 --- a/tests/common/test_element_holder.py +++ b/tests/common/test_element_holder.py @@ -50,6 +50,6 @@ def test_fill_device_delegates_to_elements(): holder = DispatchHolder() elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] - holder._fill_device(elements) + holder.fill_device(elements) assert holder.filled_elements == elements From d56d41b564076c05c60f702c42c36b3ba5dfd363 Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Mon, 7 Sep 2026 16:12:10 +0200 Subject: [PATCH 7/8] Test correction --- tests/common/test_element_holder.py | 55 ----------------------------- 1 file changed, 55 deletions(-) delete mode 100644 tests/common/test_element_holder.py diff --git a/tests/common/test_element_holder.py b/tests/common/test_element_holder.py deleted file mode 100644 index d7f66540..00000000 --- a/tests/common/test_element_holder.py +++ /dev/null @@ -1,55 +0,0 @@ -from pyaml.common.element import Element -from pyaml.common.holders.element_holder import ElementHolder - - -class DispatchElement(Element): - def _fill_device(self, holder: ElementHolder) -> None: - holder.filled_elements.append(self) - - -class DispatchHolder(ElementHolder): - def __init__(self): - super().__init__() - self.filled_elements = [] - - def create_magnet_strength_aggregator(self, magnets): - return None - - def create_magnet_hardware_aggregator(self, magnets): - return None - - def create_bpm_aggregators(self, bpms): - return [None, None, None] - - def _fill_magnet(self, magnet) -> None: - raise AssertionError("Unexpected magnet dispatch") - - def _fill_combined_function_magnet(self, magnet) -> None: - raise AssertionError("Unexpected combined-function magnet dispatch") - - def _fill_serialized_magnets(self, magnets) -> None: - raise AssertionError("Unexpected serialized magnet dispatch") - - def _fill_bpm(self, bpm) -> None: - raise AssertionError("Unexpected BPM dispatch") - - def _fill_rf_plant(self, rf_plant) -> None: - raise AssertionError("Unexpected RF plant dispatch") - - def _fill_betatron_tune_monitor(self, monitor) -> None: - raise AssertionError("Unexpected betatron tune monitor dispatch") - - def _fill_tool(self, tool) -> None: - raise AssertionError("Unexpected tool dispatch") - - def _fill_unbound_element(self, element) -> None: - raise AssertionError("Unexpected unbound element dispatch") - - -def test_fill_device_delegates_to_elements(): - holder = DispatchHolder() - elements = [DispatchElement("FIRST"), DispatchElement("SECOND")] - - holder.fill_device(elements) - - assert holder.filled_elements == elements From 9a156249a4b06eb10a3594aa26e5502b966df638 Mon Sep 17 00:00:00 2001 From: guillaumepichon Date: Tue, 8 Sep 2026 10:51:53 +0200 Subject: [PATCH 8/8] Test improvements to use existing lattices --- .../common/test_element_holder_collection.py | 196 +++++++----------- 1 file changed, 77 insertions(+), 119 deletions(-) diff --git a/tests/common/test_element_holder_collection.py b/tests/common/test_element_holder_collection.py index 7495839f..937d2da7 100644 --- a/tests/common/test_element_holder_collection.py +++ b/tests/common/test_element_holder_collection.py @@ -4,189 +4,147 @@ from pyaml.arrays.element_array import ElementArray from pyaml.arrays.magnet_array import MagnetArray from pyaml.bpm.bpm import BPM -from pyaml.common.element import Element from pyaml.common.exception import PyAMLException -from pyaml.common.holders.element_holder import ElementHolder -from pyaml.magnet.magnet import Magnet +from pyaml.lattice.simulator import Simulator -class CollectionHolder(ElementHolder): - """Minimal holder for collection examples, without device access.""" +@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 - def __init__(self, *elements): - super().__init__() - for element in elements: - element._peer = self - self.add_element(element) - def create_magnet_strength_aggregator(self, magnets): - return None - - def create_magnet_hardware_aggregator(self, magnets): - return None - - def create_bpm_aggregators(self, bpms): - return [None, None, None] - - -class Quadrupole(Magnet): - pass - - -class Corrector(Magnet): - pass - - -def test_exact_name_returns_the_element_or_none(): - bpm = BPM("BPM01") - holder = CollectionHolder(bpm) - - assert holder["BPM01"] is bpm +def test_exact_name_returns_the_element_or_none(holder): + assert holder["BPM_C04-01"] is holder.bpm.get("BPM_C04-01") assert holder["UNKNOWN"] is None -def test_get_and_full_slice_keep_registration_order(): - holder = CollectionHolder(BPM("BPM02"), BPM("BPM01")) +def test_get_and_full_slice_keep_registration_order(holder): + names = [element.get_name() for element in holder.get_all_elements()] assert type(holder.get()) is ElementArray assert type(holder[:]) is ElementArray - assert holder.get().names() == ["BPM02", "BPM01"] - assert holder[:].names() == ["BPM02", "BPM01"] + assert holder.get().names() == names + assert holder[:].names() == names + assert names != sorted(names) -def test_collections_are_independent_but_share_elements(): - bpm = BPM("BPM01") - holder = CollectionHolder(bpm) +def test_collections_are_independent_but_share_elements(holder): collection = holder.get() + first = holder[0] + original_names = collection.names() - assert collection[0] is bpm + assert collection[0] is first collection.clear() - assert holder["BPM01"] is bpm - assert holder.get().names() == ["BPM01"] + assert holder[0] is first + assert holder.get().names() == original_names -def test_new_calls_reflect_registry_additions(): - holder = CollectionHolder(BPM("BPM01")) +def test_new_calls_reflect_registry_additions(holder): previous = holder.get() - added = BPM("BPM02").attach(holder, None, None, None) - holder.add_element(added) - - assert previous.names() == ["BPM01"] - assert holder.get().names() == ["BPM01", "BPM02"] + holder.fill_device([BPM("EXTRA_BPM", lattice_names="list(BPM_C04-01)")]) + assert "EXTRA_BPM" not in previous.names() + assert holder.get().names() == previous.names() + ["EXTRA_BPM"] + assert holder["EXTRA_BPM"] is holder.bpm.get("EXTRA_BPM") -def test_patterns_always_return_arrays(): - holder = CollectionHolder(BPM("BPM02"), BPM("BPM01"), Element("MARKER")) +def test_patterns_always_return_arrays(holder): assert isinstance(holder["BPM*"], BPMArray) - assert holder["BPM*"].names() == ["BPM02", "BPM01"] - assert isinstance(holder["BPM0[1]"], BPMArray) - assert holder["BPM0[1]"].names() == ["BPM01"] - assert holder["BPM0?"].names() == ["BPM02", "BPM01"] + assert holder["BPM*"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert isinstance(holder["BPM_C04-0[1]"], BPMArray) + assert holder["BPM_C04-0[1]"].names() == ["BPM_C04-01"] + assert holder["BPM_C04-0?"].names() == ["BPM_C04-01", "BPM_C04-02"] assert type(holder["MISSING*"]) is ElementArray assert holder["MISSING*"].names() == [] -def test_character_classes_in_name_patterns(): - holder = CollectionHolder( - BPM("BPM01"), - BPM("BPM02"), - BPM("BPM03"), - BPM("BPM04"), - Magnet("QF01"), - Magnet("QD01"), - Magnet("QS01"), - ) +def test_character_classes_in_name_patterns(holder): + assert holder["BPM_C04-0[12]"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert holder["BPM_C04-0[1-2]"].names() == ["BPM_C04-01", "BPM_C04-02"] + assert holder["SH1A-C01-[HV]*"].names() == ["SH1A-C01-H", "SH1A-C01-V"] + assert holder["BPM_C04-0[!2]"].names() == ["BPM_C04-01"] - assert holder["BPM0[123]"].names() == ["BPM01", "BPM02", "BPM03"] - assert holder["BPM0[1-3]"].names() == ["BPM01", "BPM02", "BPM03"] - assert holder["Q[FD]*"].names() == ["QF01", "QD01"] - assert holder["BPM0[!3]"].names() == ["BPM01", "BPM02", "BPM04"] +def test_colons_are_part_of_names(holder): + holder.fill_device([BPM("CELL04:BPM01", lattice_names="list(BPM_C04-01)")]) -def test_colons_are_part_of_names(): - bpm = BPM("CELL01:BPM01") - holder = CollectionHolder(bpm, Magnet("MAGNET")) - - assert holder["CELL01:BPM01"] is bpm - assert holder["CELL01:BPM*"].names() == ["CELL01:BPM01"] + assert holder["CELL04:BPM01"] is holder.bpm.get("CELL04:BPM01") + assert holder["CELL04:BPM*"].names() == ["CELL04:BPM01"] assert holder["model_name:*"].names() == [] -def test_magnet_subclasses_share_a_typed_array_in_either_order(): - holder = CollectionHolder(Quadrupole("MAG_Q"), Corrector("MAG_C")) +def test_magnet_subclasses_share_a_typed_array_in_either_order(holder): + horizontal = holder.magnet.get("SH1A-C01-H") + vertical = holder.magnet.get("SH1A-C01-V") + start = holder.get_all_elements().index(horizontal) - assert isinstance(holder["MAG*"], MagnetArray) - assert isinstance(holder[::-1], MagnetArray) - assert holder[::-1].names() == ["MAG_C", "MAG_Q"] + assert type(horizontal) is not type(vertical) + assert isinstance(holder["SH1A-C01-[HV]"], MagnetArray) + assert isinstance(holder[start : start + 2], MagnetArray) + assert isinstance(holder[start + 1 : start - 1 : -1], MagnetArray) + assert holder[start + 1 : start - 1 : -1].names() == ["SH1A-C01-V", "SH1A-C01-H"] -def test_mixed_selection_returns_a_generic_array(): - holder = CollectionHolder(Magnet("CELL_Q"), BPM("CELL_BPM")) +def test_mixed_selection_returns_a_generic_array(holder): + selected = holder["*-C01*"] - assert type(holder["CELL*"]) is ElementArray - assert holder["CELL*"].names() == ["CELL_Q", "CELL_BPM"] + assert type(selected) is ElementArray + assert "QF1A-C01" in selected.names() + assert "SH1A-C01" in selected.names() -def test_indices_and_slices_follow_insertion_order(): - second = BPM("BPM02") - first = BPM("BPM01") - third = BPM("BPM03") - holder = CollectionHolder(second, first, third) +def test_indices_and_slices_follow_insertion_order(holder): + registered = holder.get_all_elements() - assert holder[0] is second - assert holder[-1] is third - assert isinstance(holder[1:3], BPMArray) - assert holder[1:3].names() == ["BPM01", "BPM03"] - assert holder[::2].names() == ["BPM02", "BPM03"] - assert type(holder[3:]) is ElementArray + assert holder[0] is registered[0] + assert holder[-1] is registered[-1] + assert list(holder[1:3]) == registered[1:3] + assert list(holder[::2]) == registered[::2] + assert isinstance(holder[-2:], BPMArray) + assert holder[-2:].names() == ["BPM_C04-01", "BPM_C04-02"] + assert type(holder[len(registered) :]) is ElementArray -def test_empty_holder_returns_empty_collections(): - holder = CollectionHolder() +def test_empty_holder_returns_empty_collections(ebs_lattice_file): + holder = Simulator(name="empty", lattice=str(ebs_lattice_file)) assert type(holder.get()) is ElementArray assert holder[:].names() == [] assert holder["BPM*"].names() == [] - assert holder["BPM01"] is None + assert holder["BPM_C04-01"] is None -def test_invalid_indices_and_keys_raise_clear_errors(): - holder = CollectionHolder(BPM("BPM01")) +def test_invalid_indices_and_keys_raise_clear_errors(holder): + size = len(holder.get_all_elements()) with pytest.raises(IndexError): - holder[1] + holder[size] with pytest.raises(IndexError): - holder[-2] + holder[-size - 1] with pytest.raises(TypeError): holder[1.5] with pytest.raises(ValueError): holder[::0] -def test_selection_intersects_with_a_configured_family(): - holder = CollectionHolder(Magnet("Q02"), Magnet("Q01"), BPM("BPM01")) - holder.fill_element_array("CELL01", ["Q01", "BPM01"]) - - selected = holder["Q*"] & holder.get_elements("CELL01") +def test_selection_intersects_with_a_configured_family(holder): + selected = holder["SH1A-C0?-H"] & holder.get_elements("ElArray") assert isinstance(selected, MagnetArray) - assert selected.names() == ["Q01"] - assert holder.get_element("Q01") is holder["Q01"] + assert selected.names() == ["SH1A-C02-H"] + assert holder.get_element("SH1A-C02-H") is holder["SH1A-C02-H"] assert holder.get_all_elements() == list(holder.get()) with pytest.raises(PyAMLException): holder.get_element("UNKNOWN") -def test_existing_array_field_filters_still_work(): - magnet = Magnet("VIRTUAL_Q") - magnet.set_model_name("PHYSICAL_Q") - holder = CollectionHolder(magnet) - - assert holder["VIRTUAL*"]["model_name:PHYSICAL*"].names() == ["VIRTUAL_Q"] +def test_existing_array_field_filters_still_work(holder): + selected = holder["SH1A-C0?-H"]["model_name:SH1A-C01"] + assert selected.names() == ["SH1A-C01-H"] -def test_existing_empty_intersection_stays_a_list(): - holder = CollectionHolder(Magnet("Q01"), Magnet("Q02")) - assert type(holder["Q0[1]"] & holder["Q0[2]"]) is list +def test_existing_empty_intersection_stays_a_list(holder): + assert type(holder["SH1A-C0?-H"] & holder["SH1A-C0?-V"]) is list