add label based indexing to pass sequence (#154)

* add label based indexing to pass sequence

* add ipython completion support

Author: axtimhaus

pyroll/core/sequence/sequence.py

@@ -58,15 +58,31 @@ def __iter__(self):
         return self._subunits.__iter__()
 
     @overload
-    def __getitem__(self, index: int) -> Unit:
+    def __getitem__(self, key: int) -> Unit:
+        """Gets unit item by index."""
         ...
 
     @overload
-    def __getitem__(self, index: slice) -> list[Unit]:
+    def __getitem__(self, key: str) -> Unit:
+        """Gets unit item by label."""
         ...
 
-    def __getitem__(self, index: int) -> Unit:
-        return self._subunits.__getitem__(index)
+    @overload
+    def __getitem__(self, key: slice) -> list[Unit]:
+        """Gets a slice of units."""
+        ...
+
+    def __getitem__(self, key):
+        if isinstance(key, str):
+            try:
+                return next(u for u in self._subunits if u.label == key)
+            except StopIteration:
+                raise KeyError(f"No unit with label '{key}' found.")
+
+        if isinstance(key, int) or isinstance(key, slice):
+            return self._subunits.__getitem__(key)
+
+        raise TypeError("Key must be int, slice or str")
 
     @property
     def units(self) -> List[Unit]:
@@ -88,3 +104,6 @@ def __attrs__(self):
         return super().__attrs__ | {
             "units": self.units
         }
+    
+    def _ipython_key_completions_(self):
+        return [u.label for u in self._subunits]

tests/test_pass_sequence.py

@@ -0,0 +1,63 @@
+from pyroll.core import PassSequence, RollPass, CircularOvalGroove, Transport, RoundGroove, Roll
+import pytest
+
+
+# noinspection DuplicatedCode
+def test_pass_sequence_indexing():
+    sequence = PassSequence([
+        RollPass(
+            label="Oval I",
+            roll=Roll(
+                groove=CircularOvalGroove(
+                    depth=8e-3,
+                    r1=6e-3,
+                    r2=40e-3
+                ),
+                nominal_radius=160e-3,
+                rotational_frequency=1,
+                neutral_point=-20e-3
+            ),
+            gap=2e-3,
+
+        ),
+        Transport(
+            label="I => II",
+            duration=1,
+        ),
+        RollPass(
+            label="Round II",
+            roll=Roll(
+                groove=RoundGroove(
+                    r1=1e-3,
+                    r2=12.5e-3,
+                    depth=11.5e-3
+                ),
+                nominal_radius=160e-3,
+                rotational_frequency=1
+            ),
+            gap=2e-3,
+        ),
+        Transport(
+            label="II => III",
+            duration=1
+        ),
+        RollPass(
+            label="Oval III",
+            roll=Roll(
+                groove=CircularOvalGroove(
+                    depth=6e-3,
+                    r1=6e-3,
+                    r2=35e-3
+                ),
+                nominal_radius=160e-3,
+                rotational_frequency=1
+            ),
+            gap=2e-3,
+        ),
+    ])
+
+    assert sequence["Oval III"] == sequence[4]
+    assert sequence["I => II"] == sequence[1]
+
+    with pytest.raises(KeyError):
+        _ = sequence["not present"]

tests/test_solve.py

@@ -11,6 +11,7 @@ def flow_stress(self: RollPass.Profile):
     return 50e6 * (1 + self.strain) ** 0.2 * self.roll_pass.strain_rate ** 0.1
 
 
+# noinspection DuplicatedCode
 def test_solve(tmp_path: Path, caplog):
     caplog.set_level(logging.DEBUG, logger="pyroll")