Skip to content

Commit e08e762

Browse files
Merge pull request #376 from python-accelerator-middle-layer/366-feature-add-generic-elementholder-collection-access
feat: add generic ElementHolder collection access
2 parents dcf2c8e + 7fdb3ce commit e08e762

4 files changed

Lines changed: 280 additions & 1 deletion

File tree

‎pyaml/arrays/element_array.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,25 @@ def __auto_array(self, elements: list[Element]):
182182
if len(elements) == 0:
183183
return []
184184

185+
return self._typed_array(elements)
186+
187+
def _typed_array(self, elements: list[Element]) -> "ElementArray":
188+
"""Build a collection using the most specific compatible array type.
189+
190+
Parameters
191+
----------
192+
elements : list[Element]
193+
Selected references, in their desired order.
194+
195+
Returns
196+
-------
197+
ElementArray
198+
Specialized array when possible, otherwise a generic array.
199+
An empty selection returns an empty generic array.
200+
"""
201+
if not elements:
202+
return self.__create_array("", Element, elements)
203+
185204
import inspect
186205

187206
def mro_as_list(cls: type) -> list[type]:
@@ -210,6 +229,22 @@ def mro_as_list(cls: type) -> list[type]:
210229

211230
return self.__create_array("", chosen, elements)
212231

232+
def _select_names(self, pattern: str) -> "ElementArray":
233+
"""Select names without interpreting field selectors.
234+
235+
Parameters
236+
----------
237+
pattern : str
238+
A fnmatch pattern applied to each element name.
239+
240+
Returns
241+
-------
242+
ElementArray
243+
Typed selection in the original order, including an empty array
244+
when no names match.
245+
"""
246+
return self._typed_array([element for element in self if fnmatch.fnmatch(element.get_name(), pattern)])
247+
213248
def __is_bool_mask(self, other: object) -> bool:
214249
"""Return True if 'other' looks like a boolean mask (list or numpy array)."""
215250
# --- numpy boolean array ---

‎pyaml/common/holders/element_holder.py‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import fnmatch
44
import re
55
from abc import ABCMeta, abstractmethod
6-
from typing import TYPE_CHECKING
6+
from typing import TYPE_CHECKING, overload
77

88
from ...arrays.element_array import ElementArray
99
from ...bpm.bpm import BPM
@@ -419,6 +419,95 @@ def _get(self, what, name, array) -> Element:
419419
return array[name]
420420

421421
# Generic elements
422+
def get(self) -> ElementArray:
423+
"""Return all registered elements in insertion order.
424+
425+
Returns
426+
-------
427+
ElementArray
428+
New unnamed container sharing the registered element references.
429+
430+
Notes
431+
-----
432+
Registration order is not necessarily longitudinal lattice order.
433+
Changing the returned container does not change the holder registry.
434+
Each call reflects the current registry.
435+
436+
Examples
437+
--------
438+
>>> elements = sr.live.get()
439+
>>> names = elements.names()
440+
"""
441+
return ElementArray("", list(self._ALL.values()))
442+
443+
@overload
444+
def __getitem__(self, key: int) -> Element: ...
445+
446+
@overload
447+
def __getitem__(self, key: slice) -> ElementArray: ...
448+
449+
@overload
450+
def __getitem__(self, key: str) -> Element | ElementArray | None: ...
451+
452+
def __getitem__(self, key: int | slice | str) -> Element | ElementArray | None:
453+
"""Retrieve an element or select a collection.
454+
455+
Parameters
456+
----------
457+
key : int, slice or str
458+
Index in registration order, slice, exact name, or name pattern.
459+
Strings containing ``*``, ``?`` or ``[`` use fnmatch matching.
460+
Other strings are exact registry keys. Colons are literal.
461+
462+
Returns
463+
-------
464+
Element or ElementArray or None
465+
An index returns an element. An exact name returns its element
466+
or None. Patterns and slices return the most specific compatible
467+
array, or an empty ElementArray when nothing matches.
468+
The full slice ``[:]`` returns a generic ElementArray, like get().
469+
470+
Raises
471+
------
472+
IndexError
473+
If the index is out of bounds.
474+
TypeError
475+
If the key is neither an integer, a slice, nor a string.
476+
ValueError
477+
If a slice has a zero step.
478+
479+
Notes
480+
-----
481+
Indices follow insertion order, not necessarily lattice order.
482+
Collections share element references but do not modify the registry.
483+
Field filters and regular expressions are not interpreted here.
484+
485+
Examples
486+
--------
487+
>>> bpm = sr.live["BPM01"]
488+
>>> missing = sr.live["UNKNOWN"] # None
489+
>>> bpms = sr.live["BPM*"]
490+
>>> bpms = sr.live["BPM0[123]"] # BPM01, BPM02 or BPM03
491+
>>> bpms = sr.live["BPM0[1-3]"] # Same selection using a range
492+
>>> quads = sr.live["Q[FD]*"] # Names starting with QF or QD
493+
>>> bpms = sr.live["BPM0[!3]"] # One character after BPM0, except 3
494+
>>> first = sr.live[0]
495+
>>> subset = sr.live[1:10]
496+
>>> all_elements = sr.live[:]
497+
"""
498+
if isinstance(key, str):
499+
if any(marker in key for marker in "*?["):
500+
return self.get()._select_names(key)
501+
return self._ALL.get(key)
502+
if isinstance(key, int):
503+
return list(self._ALL.values())[key]
504+
if isinstance(key, slice):
505+
elements = self.get()
506+
if key == slice(None):
507+
return elements
508+
return elements._typed_array(list(elements)[key])
509+
raise TypeError("ElementHolder keys must be integers, slices or strings")
510+
422511
def fill_element_array(self, arrayName: str, elementNames: list[str]):
423512
"""
424513
Create and register a generic element array.
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import pytest
2+
3+
from pyaml.arrays.bpm_array import BPMArray
4+
from pyaml.arrays.element_array import ElementArray
5+
from pyaml.arrays.magnet_array import MagnetArray
6+
from pyaml.bpm.bpm import BPM
7+
from pyaml.common.exception import PyAMLException
8+
from pyaml.lattice.simulator import Simulator
9+
10+
11+
@pytest.fixture
12+
def holder(accelerator_from_fragments, sr_configuration_fragments):
13+
sr = accelerator_from_fragments(*sr_configuration_fragments)
14+
sr.design.get_lattice().disable_6d()
15+
return sr.design
16+
17+
18+
def test_exact_name_returns_the_element_or_none(holder):
19+
assert holder["BPM_C04-01"] is holder.bpm.get("BPM_C04-01")
20+
assert holder["UNKNOWN"] is None
21+
22+
23+
def test_get_and_full_slice_keep_registration_order(holder):
24+
names = [element.get_name() for element in holder.get_all_elements()]
25+
26+
assert type(holder.get()) is ElementArray
27+
assert type(holder[:]) is ElementArray
28+
assert holder.get().names() == names
29+
assert holder[:].names() == names
30+
assert names != sorted(names)
31+
32+
33+
def test_collections_are_independent_but_share_elements(holder):
34+
collection = holder.get()
35+
first = holder[0]
36+
original_names = collection.names()
37+
38+
assert collection[0] is first
39+
collection.clear()
40+
assert holder[0] is first
41+
assert holder.get().names() == original_names
42+
43+
44+
def test_new_calls_reflect_registry_additions(holder):
45+
previous = holder.get()
46+
holder.fill_device([BPM("EXTRA_BPM", lattice_names="list(BPM_C04-01)")])
47+
48+
assert "EXTRA_BPM" not in previous.names()
49+
assert holder.get().names() == previous.names() + ["EXTRA_BPM"]
50+
assert holder["EXTRA_BPM"] is holder.bpm.get("EXTRA_BPM")
51+
52+
53+
def test_patterns_always_return_arrays(holder):
54+
assert isinstance(holder["BPM*"], BPMArray)
55+
assert holder["BPM*"].names() == ["BPM_C04-01", "BPM_C04-02"]
56+
assert isinstance(holder["BPM_C04-0[1]"], BPMArray)
57+
assert holder["BPM_C04-0[1]"].names() == ["BPM_C04-01"]
58+
assert holder["BPM_C04-0?"].names() == ["BPM_C04-01", "BPM_C04-02"]
59+
assert type(holder["MISSING*"]) is ElementArray
60+
assert holder["MISSING*"].names() == []
61+
62+
63+
def test_character_classes_in_name_patterns(holder):
64+
assert holder["BPM_C04-0[12]"].names() == ["BPM_C04-01", "BPM_C04-02"]
65+
assert holder["BPM_C04-0[1-2]"].names() == ["BPM_C04-01", "BPM_C04-02"]
66+
assert holder["SH1A-C01-[HV]*"].names() == ["SH1A-C01-H", "SH1A-C01-V"]
67+
assert holder["BPM_C04-0[!2]"].names() == ["BPM_C04-01"]
68+
69+
70+
def test_colons_are_part_of_names(holder):
71+
holder.fill_device([BPM("CELL04:BPM01", lattice_names="list(BPM_C04-01)")])
72+
73+
assert holder["CELL04:BPM01"] is holder.bpm.get("CELL04:BPM01")
74+
assert holder["CELL04:BPM*"].names() == ["CELL04:BPM01"]
75+
assert holder["model_name:*"].names() == []
76+
77+
78+
def test_magnet_subclasses_share_a_typed_array_in_either_order(holder):
79+
horizontal = holder.magnet.get("SH1A-C01-H")
80+
vertical = holder.magnet.get("SH1A-C01-V")
81+
start = holder.get_all_elements().index(horizontal)
82+
83+
assert type(horizontal) is not type(vertical)
84+
assert isinstance(holder["SH1A-C01-[HV]"], MagnetArray)
85+
assert isinstance(holder[start : start + 2], MagnetArray)
86+
assert isinstance(holder[start + 1 : start - 1 : -1], MagnetArray)
87+
assert holder[start + 1 : start - 1 : -1].names() == ["SH1A-C01-V", "SH1A-C01-H"]
88+
89+
90+
def test_mixed_selection_returns_a_generic_array(holder):
91+
selected = holder["*-C01*"]
92+
93+
assert type(selected) is ElementArray
94+
assert "QF1A-C01" in selected.names()
95+
assert "SH1A-C01" in selected.names()
96+
97+
98+
def test_indices_and_slices_follow_insertion_order(holder):
99+
registered = holder.get_all_elements()
100+
101+
assert holder[0] is registered[0]
102+
assert holder[-1] is registered[-1]
103+
assert list(holder[1:3]) == registered[1:3]
104+
assert list(holder[::2]) == registered[::2]
105+
assert isinstance(holder[-2:], BPMArray)
106+
assert holder[-2:].names() == ["BPM_C04-01", "BPM_C04-02"]
107+
assert type(holder[len(registered) :]) is ElementArray
108+
109+
110+
def test_empty_holder_returns_empty_collections(ebs_lattice_file):
111+
holder = Simulator(name="empty", lattice=str(ebs_lattice_file))
112+
113+
assert type(holder.get()) is ElementArray
114+
assert holder[:].names() == []
115+
assert holder["BPM*"].names() == []
116+
assert holder["BPM_C04-01"] is None
117+
118+
119+
def test_invalid_indices_and_keys_raise_clear_errors(holder):
120+
size = len(holder.get_all_elements())
121+
122+
with pytest.raises(IndexError):
123+
holder[size]
124+
with pytest.raises(IndexError):
125+
holder[-size - 1]
126+
with pytest.raises(TypeError):
127+
holder[1.5]
128+
with pytest.raises(ValueError):
129+
holder[::0]
130+
131+
132+
def test_selection_intersects_with_a_configured_family(holder):
133+
selected = holder["SH1A-C0?-H"] & holder.get_elements("ElArray")
134+
135+
assert isinstance(selected, MagnetArray)
136+
assert selected.names() == ["SH1A-C02-H"]
137+
assert holder.get_element("SH1A-C02-H") is holder["SH1A-C02-H"]
138+
assert holder.get_all_elements() == list(holder.get())
139+
with pytest.raises(PyAMLException):
140+
holder.get_element("UNKNOWN")
141+
142+
143+
def test_existing_array_field_filters_still_work(holder):
144+
selected = holder["SH1A-C0?-H"]["model_name:SH1A-C01"]
145+
146+
assert selected.names() == ["SH1A-C01-H"]
147+
148+
149+
def test_existing_empty_intersection_stays_a_list(holder):
150+
assert type(holder["SH1A-C0?-H"] & holder["SH1A-C0?-V"]) is list

‎tests/test_load_conf_with_code.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,8 @@ def test_load_conf_with_code():
1111
bpms = sr.live.bpms.get("BPM")
1212
assert bpms is not None
1313
assert len(bpms) == 320
14+
15+
assert sr.live[bpms[0].get_name()] is bpms[0]
16+
assert sr.live["BPM*"].names() == bpms.names()
17+
assert sr.live[:].names() == [element.get_name() for element in sr.live.get_all_elements()]
18+
assert sr.design["BPM*"].names() == sr.design.bpms.get("BPM").names()

0 commit comments

Comments
 (0)