Implement UnionEle

Deduplicate logic with `BeamLine` for allowed types
and (de)serialization.

Author: ax3l

src/pals/kinds/BeamLine.py

@@ -1,156 +1,27 @@
-from pydantic import ConfigDict, Field, model_validator
-from typing import Annotated, List, Literal, Union
+from pydantic import model_validator
+from typing import List, Literal
 
-from .BaseElement import BaseElement
-from .ThickElement import ThickElement
-
-from .ACKicker import ACKicker
-from .BeamBeam import BeamBeam
-from .BeginningEle import BeginningEle
-from .Converter import Converter
-from .CrabCavity import CrabCavity
-from .Drift import Drift
-from .EGun import EGun
-from .Feedback import Feedback
-from .Fiducial import Fiducial
-from .FloorShift import FloorShift
-from .Foil import Foil
-from .Fork import Fork
-from .Girder import Girder
-from .Instrument import Instrument
-from .Kicker import Kicker
-from .Marker import Marker
-from .Mask import Mask
-from .Match import Match
-from .Multipole import Multipole
-from .NullEle import NullEle
-from .Octupole import Octupole
-from .Patch import Patch
-from .Quadrupole import Quadrupole
-from .RBend import RBend
-from .RFCavity import RFCavity
-from .SBend import SBend
-from .Sextupole import Sextupole
-from .Solenoid import Solenoid
-from .Taylor import Taylor
-from .UnionEle import UnionEle
-from .Wiggler import Wiggler
+from .all_elements import get_all_elements_as_annotation
+from .mixin import BaseElement
 
 
 class BeamLine(BaseElement):
     """A line of elements and/or other lines"""
 
-    # Validate every time a new value is assigned to an attribute,
-    # not only when an instance of BeamLine is created
-    model_config = ConfigDict(validate_assignment=True)
-
     kind: Literal["BeamLine"] = "BeamLine"
 
-    line: List[
-        Annotated[
-            Union[
-                # Base classes (for testing compatibility)
-                BaseElement,
-                ThickElement,
-                # User-Facing element kinds
-                "BeamLine",
-                ACKicker,
-                BeamBeam,
-                BeginningEle,
-                Converter,
-                CrabCavity,
-                Drift,
-                EGun,
-                Feedback,
-                Fiducial,
-                FloorShift,
-                Foil,
-                Fork,
-                Girder,
-                Instrument,
-                Kicker,
-                Marker,
-                Mask,
-                Match,
-                Multipole,
-                NullEle,
-                Octupole,
-                Patch,
-                Quadrupole,
-                RBend,
-                RFCavity,
-                SBend,
-                Sextupole,
-                Solenoid,
-                Taylor,
-                UnionEle,
-                Wiggler,
-            ],
-            Field(discriminator="kind"),
-        ]
-    ]
+    line: List[get_all_elements_as_annotation()]
 
     @model_validator(mode="before")
     @classmethod
     def unpack_yaml_structure(cls, data):
-        # Handle the top-level one-key dict: unpack the line's name
-        if isinstance(data, dict) and len(data) == 1:
-            name, value = list(data.items())[0]
-            if not isinstance(value, dict):
-                raise TypeError(
-                    f"Value for line key {name!r} must be a dict, but we got {value!r}"
-                )
-            value["name"] = name
-            data = value
-        # Handle the 'line' field: unpack each element's name
-        if "line" not in data:
-            raise ValueError("'line' field is missing")
-        if not isinstance(data["line"], list):
-            raise TypeError("'line' must be a list")
-        new_line = []
-        # Loop over all elements in the line
-        for item in data["line"]:
-            # An element can be a string that refers to another element
-            if isinstance(item, str):
-                raise RuntimeError("Reference/alias elements not yet implemented")
-            # An element can be a dict
-            elif isinstance(item, dict):
-                if not (len(item) == 1):
-                    raise ValueError(
-                        f"Each element must be a dict with exactly one key (the element's name), but we got {item!r}"
-                    )
-                name, fields = list(item.items())[0]
-                if not isinstance(fields, dict):
-                    raise TypeError(
-                        f"Value for element key {name!r} must be a dict (the element's properties), but we got {fields!r}"
-                    )
-                fields["name"] = name
-                new_line.append(fields)
-            # An element can be an instance of an existing model
-            elif isinstance(item, BaseElement):
-                # Nothing to do, keep the element as is
-                new_line.append(item)
-            else:
-                raise TypeError(
-                    f"Value for element key {name!r} must be a reference string or a dict, but we got {item!r}"
-                )
-        data["line"] = new_line
-        return data
+        """Unpack YAML/JSON/...-like structure for BeamLine elements"""
+        from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
 
-    def model_dump(self, *args, **kwargs):
-        """This makes sure the element name property is moved out and up to a one-key dictionary"""
-        # Use base element dump first and return a dict {key: value}, where 'key'
-        # is the name of the line and 'value' is a dict with all other properties
-        data = super().model_dump(*args, **kwargs)
-        # Reformat 'line' field as list of element dicts
-        new_line = []
-        for elem in self.line:
-            # Use custom dump for each line element, which now returns a dict
-            elem_dict = elem.model_dump(**kwargs)
-            new_line.append(elem_dict)
-        data[self.name]["line"] = new_line
-        return data
+        return unpack_element_list_structure(data, "line", "line")
 
+    def model_dump(self, *args, **kwargs):
+        """Custom model dump for BeamLine to handle element list formatting"""
+        from pals.kinds.mixin.all_element_mixin import dump_element_list
 
-# Avoid circular import issues
-BeamLine.model_rebuild()
+        return dump_element_list(self, "line", *args, **kwargs)

src/pals/kinds/UnionEle.py

@@ -1,4 +1,4 @@
-from pydantic import model_validator
+from pydantic import model_validator  # noqa
 from typing import List, Literal
 
 from .all_elements import get_all_elements_as_annotation
@@ -11,6 +11,19 @@ class UnionEle(BaseElement):
     # Discriminator field
     kind: Literal["UnionEle"] = "UnionEle"
 
-    # Elements in the union
-    #   Note: https://github.com/campa-consortium/pals/issues/89
-    elements: List[BaseElement] = []
+    # Elements in the union - uses the same union type as BeamLine
+    elements: List[get_all_elements_as_annotation()] = []
+
+    @model_validator(mode="before")
+    @classmethod
+    def unpack_yaml_structure(cls, data):
+        """Unpack YAML/JSON/...-like structure for UnionEle elements"""
+        from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
+
+        return unpack_element_list_structure(data, "elements", "union")
+
+    def model_dump(self, *args, **kwargs):
+        """Custom model dump for UnionEle to handle element list formatting"""
+        from pals.kinds.mixin.all_element_mixin import dump_element_list
+
+        return dump_element_list(self, "elements", *args, **kwargs)

src/pals/kinds/__init__.py

@@ -34,3 +34,8 @@
 from .Taylor import Taylor  # noqa: F401
 from .UnionEle import UnionEle  # noqa: F401
 from .Wiggler import Wiggler  # noqa: F401
+
+
+# Rebuild pydantic models that depend on other classes
+BeamLine.model_rebuild()
+UnionEle.model_rebuild()

src/pals/kinds/all_elements.py

@@ -0,0 +1,87 @@
+"""Helper module to define the union of all allowed element types.
+
+This module creates a helper function that returns the element union type,
+avoiding duplication between BeamLine.line and UnionEle.elements.
+"""
+
+from typing import Annotated, Union
+
+from pydantic import Field
+
+from .ACKicker import ACKicker
+from .BeamBeam import BeamBeam
+from .BeginningEle import BeginningEle
+from .Converter import Converter
+from .CrabCavity import CrabCavity
+from .Drift import Drift
+from .EGun import EGun
+from .Feedback import Feedback
+from .Fiducial import Fiducial
+from .FloorShift import FloorShift
+from .Foil import Foil
+from .Fork import Fork
+from .Girder import Girder
+from .Instrument import Instrument
+from .Kicker import Kicker
+from .Marker import Marker
+from .Mask import Mask
+from .Match import Match
+from .Multipole import Multipole
+from .NullEle import NullEle
+from .Octupole import Octupole
+from .Patch import Patch
+from .Quadrupole import Quadrupole
+from .RBend import RBend
+from .RFCavity import RFCavity
+from .SBend import SBend
+from .Sextupole import Sextupole
+from .Solenoid import Solenoid
+from .Taylor import Taylor
+from .Wiggler import Wiggler
+
+
+def get_all_element_types(extra_types: tuple = None):
+    """Return a tuple of all element types that can be used in BeamLine or UnionEle."""
+    element_types = (
+        "BeamLine",  # Forward reference to handle circular import
+        "UnionEle",  # Forward reference to handle circular import
+        ACKicker,
+        BeamBeam,
+        BeginningEle,
+        Converter,
+        CrabCavity,
+        Drift,
+        EGun,
+        Feedback,
+        Fiducial,
+        FloorShift,
+        Foil,
+        Fork,
+        Girder,
+        Instrument,
+        Kicker,
+        Marker,
+        Mask,
+        Match,
+        Multipole,
+        NullEle,
+        Octupole,
+        Patch,
+        Quadrupole,
+        RBend,
+        RFCavity,
+        SBend,
+        Sextupole,
+        Solenoid,
+        Taylor,
+        Wiggler,
+    )
+    if extra_types is not None:
+        element_types += extra_types
+    return element_types
+
+
+def get_all_elements_as_annotation(extra_types: tuple = None):
+    """Return the Union type of all allowed elements with their name as the discriminator field."""
+    types = get_all_element_types(extra_types)
+    return Annotated[Union[types], Field(discriminator="kind")]

tests/test_elements.py

@@ -466,7 +466,17 @@ def test_Foil():
 
 def test_UnionEle():
     """Test UnionEle element"""
+    # Test empty union
     element = pals.UnionEle(name="union1", elements=[])
     assert element.name == "union1"
     assert element.kind == "UnionEle"
     assert element.elements == []
+
+    # Test union with elements
+    marker = pals.Marker(name="m1")
+    drift = pals.Drift(name="d1", length=1.0)
+    element_with_children = pals.UnionEle(name="union2", elements=[marker, drift])
+    assert element_with_children.name == "union2"
+    assert len(element_with_children.elements) == 2
+    assert element_with_children.elements[0].name == "m1"
+    assert element_with_children.elements[1].name == "d1"

tests/test_serialization.py

@@ -178,8 +178,10 @@ def test_comprehensive_lattice():
     # Taylor
     taylor = pals.Taylor(name="taylor1")
 
-    # UnionEle
-    unionele = pals.UnionEle(name="unionele1", elements=[])
+    # UnionEle - with nested elements
+    union_marker = pals.Marker(name="union_marker")
+    union_drift = pals.Drift(name="union_drift", length=0.1)
+    unionele = pals.UnionEle(name="unionele1", elements=[union_marker, union_drift])
 
     # Wiggler
     wiggler = pals.Wiggler(name="wiggler1", length=2.0)
@@ -246,6 +248,7 @@ def test_comprehensive_lattice():
     octupole_loaded = None
     rbend_loaded = None
     rfcavity_loaded = None
+    unionele_loaded = None
 
     for elem in loaded_lattice.line:
         if elem.name == "sextupole1":
@@ -256,6 +259,8 @@ def test_comprehensive_lattice():
             rbend_loaded = elem
         elif elem.name == "rfcavity1":
             rfcavity_loaded = elem
+        elif elem.name == "unionele1":
+            unionele_loaded = elem
 
     # Test that parameter groups are correctly deserialized
     assert sextupole_loaded.MagneticMultipoleP.Bn2 == 1.0
@@ -270,6 +275,15 @@ def test_comprehensive_lattice():
     assert rfcavity_loaded.RFP.frequency == 1e9
     assert rfcavity_loaded.SolenoidP.Ksol == 0.05
 
+    # Test that UnionEle elements are correctly deserialized
+    assert unionele_loaded is not None
+    assert len(unionele_loaded.elements) == 2
+    assert unionele_loaded.elements[0].name == "union_marker"
+    assert unionele_loaded.elements[0].kind == "Marker"
+    assert unionele_loaded.elements[1].name == "union_drift"
+    assert unionele_loaded.elements[1].kind == "Drift"
+    assert unionele_loaded.elements[1].length == 0.1
+
     # Test serialization to JSON
     json_data = json.dumps(lattice.model_dump(), sort_keys=True, indent=2)
     print(f"\nComprehensive lattice JSON:\n{json_data}")
@@ -294,6 +308,7 @@ def test_comprehensive_lattice():
     octupole_loaded_json = None
     rbend_loaded_json = None
     rfcavity_loaded_json = None
+    unionele_loaded_json = None
 
     for elem in loaded_lattice_json.line:
         if elem.name == "sextupole1":
@@ -304,6 +319,8 @@ def test_comprehensive_lattice():
             rbend_loaded_json = elem
         elif elem.name == "rfcavity1":
             rfcavity_loaded_json = elem
+        elif elem.name == "unionele1":
+            unionele_loaded_json = elem
 
     # Test that parameter groups are correctly deserialized
     assert sextupole_loaded_json.MagneticMultipoleP.Bn2 == 1.0
@@ -318,6 +335,15 @@ def test_comprehensive_lattice():
     assert rfcavity_loaded_json.RFP.frequency == 1e9
     assert rfcavity_loaded_json.SolenoidP.Ksol == 0.05
 
+    # Test that UnionEle elements are correctly deserialized from JSON
+    assert unionele_loaded_json is not None
+    assert len(unionele_loaded_json.elements) == 2
+    assert unionele_loaded_json.elements[0].name == "union_marker"
+    assert unionele_loaded_json.elements[0].kind == "Marker"
+    assert unionele_loaded_json.elements[1].name == "union_drift"
+    assert unionele_loaded_json.elements[1].kind == "Drift"
+    assert unionele_loaded_json.elements[1].length == 0.1
+
     # Clean up temporary files
     os.remove(yaml_file)
     os.remove(json_file)