Add ElementReference

Used inside `BeamLine::line` and later on similarly used
for `Lattice::branches` and other beamline extensions.

Author: ax3l

src/pals/kinds/ElementReference.py

@@ -0,0 +1,83 @@
+"""Element reference class for referring to elements by name."""
+
+from pydantic import BaseModel, Field, model_serializer
+from typing import Annotated
+
+from .mixin import BaseElement
+
+
+class ElementReference(BaseModel):
+    """A pydantic model that represents a reference to a named element.
+
+    This class behaves like a string (via __str__ and __eq__) but stores
+    a true reference to the actual element object once it's resolved.
+
+    The element field holds a reference (not a copy) to the actual element.
+
+    Attributes:
+        name: The name of the referenced element
+        element: A reference to the resolved element object (None until resolved)
+
+    Example:
+        >>> ref = ElementReference(name="drift1")
+        >>> ref.name
+        'drift1'
+        >>> str(ref)
+        'drift1'
+        >>> ref == "drift1"
+        True
+        >>> ref.element  # None until resolved
+        >>> drift = pals.Drift(name="drift1", length=1.0)
+        >>> ref.element = drift
+        >>> ref.is_resolved()
+        True
+        >>> ref.element is drift  # True - it's a reference, not a copy
+        True
+    """
+
+    name: str = Field(..., description="The name of the referenced element")
+    element: Annotated[
+        "BaseElement | None",
+        Field(default=None, description="Reference to the resolved element object"),
+    ] = None
+
+    @model_serializer(mode="plain")
+    def _serialize_as_name(self) -> str:
+        """Serialize this reference as just its name.
+
+        This makes `model_dump()` return a string (the element name), so nested
+        serialization (e.g. inside BeamLine.line) produces plain strings too.
+        """
+        return self.name
+
+    def __init__(self, name: str | None = None, /, **data):
+        """Initialize with either positional name or keyword arguments."""
+        if name is not None:
+            super().__init__(name=name, **data)
+        else:
+            super().__init__(**data)
+
+    def __str__(self) -> str:
+        """Return the element name as string."""
+        return self.name
+
+    def __eq__(self, other: object) -> bool:
+        """Enable string comparison."""
+        if isinstance(other, str):
+            return self.name == other
+        if isinstance(other, ElementReference):
+            return self.name == other.name and self.element is other.element
+        return False
+
+    def __hash__(self) -> int:
+        """Make hashable like a string."""
+        return hash(self.name)
+
+    def is_resolved(self) -> bool:
+        """Check if this reference has been resolved to an actual element."""
+        return self.element is not None
+
+    def __repr__(self) -> str:
+        """Return a representation of the ElementReference."""
+        resolved = "resolved" if self.is_resolved() else "unresolved"
+        return f"ElementReference('{self.name}', {resolved})"

src/pals/kinds/__init__.py

@@ -10,6 +10,7 @@
 from .CrabCavity import CrabCavity  # noqa: F401
 from .Drift import Drift  # noqa: F401
 from .EGun import EGun  # noqa: F401
+from .ElementReference import ElementReference  # noqa: F401
 from .Feedback import Feedback  # noqa: F401
 from .Fiducial import Fiducial  # noqa: F401
 from .FloorShift import FloorShift  # noqa: F401

src/pals/kinds/all_elements.py

@@ -4,9 +4,8 @@
 avoiding duplication between BeamLine.line and UnionEle.elements.
 """
 
-from typing import Annotated, Union
+from typing import Union
 
-from pydantic import Field
 
 from .ACKicker import ACKicker
 from .BeamBeam import BeamBeam
@@ -15,6 +14,7 @@
 from .CrabCavity import CrabCavity
 from .Drift import Drift
 from .EGun import EGun
+from .ElementReference import ElementReference
 from .Feedback import Feedback
 from .Fiducial import Fiducial
 from .FloorShift import FloorShift
@@ -84,11 +84,11 @@ def get_all_element_types(extra_types: tuple = None):
 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.
 
-    Note: When str is included in the union (for string references), we cannot use
-    discriminator since str doesn't have a 'kind' field. Pydantic will still properly
-    validate the union by trying each type in order.
+    Note: ElementReference is included to support string references to named elements.
+    Since ElementReference doesn't have a 'kind' field, we cannot use discriminator.
+    Pydantic will still properly validate the union by trying each type in order in
+    our unpack_element_list_structure method.
     """
-    types = get_all_element_types(extra_types)
-    # Add str to support string references to named elements
-    # We can't use discriminator with str in the union since str has no 'kind' field
-    return Union[types + (str,)]
+    types = get_all_element_types(extra_types) + (ElementReference,)
+    # We can't use discriminator with ElementReference in the union since it has no 'kind' field
+    return Union[types]

src/pals/kinds/mixin/all_element_mixin.py

@@ -5,6 +5,7 @@
 """
 
 from . import BaseElement
+from ..ElementReference import ElementReference
 
 
 def unpack_element_list_structure(
@@ -44,7 +45,12 @@ def unpack_element_list_structure(
     for item in data[field_name]:
         # An element can be a string that refers to another element
         if isinstance(item, str):
-            # Keep the string reference as-is - it will be validated later
+            # Wrap the string in an ElementReference object
+            new_list.append(ElementReference(item))
+            continue
+        # An element can be an ElementReference instance directly
+        elif isinstance(item, ElementReference):
+            # Keep the ElementReference as-is
             new_list.append(item)
             continue
         # An element can be a dict
@@ -71,7 +77,7 @@ def unpack_element_list_structure(
                 continue
 
             raise TypeError(
-                f"Value must be a reference string or a dict, but we got {item!r}"
+                f"Value must be a reference string, ElementReference, or a dict, but we got {item!r}"
             )
 
     data[field_name] = new_list

src/pals/schema_version.py

@@ -0,0 +1,4 @@
+from typing import Optional
+
+# PALS schema version - null for now, will be set when version scheme is finalized
+PALS_SCHEMA_VERSION: Optional[str] = None

tests/test_elements.py

@@ -534,11 +534,67 @@ def test_BeamLine_with_string_references():
 
     assert beamline.name == "fodo_cell"
     assert len(beamline.line) == 3
-    # First element should be a string reference
+
+    # First element should be an ElementReference that behaves like the string "drift1"
+    assert isinstance(beamline.line[0], pals.ElementReference)
     assert beamline.line[0] == "drift1"
-    # Second element should be a string reference
+    assert beamline.line[0].name == "drift1"
+    assert beamline.line[0].element is None  # Not yet resolved
+    assert not beamline.line[0].is_resolved()
+
+    # Second element should be an ElementReference that behaves like the string "quad1"
+    assert isinstance(beamline.line[1], pals.ElementReference)
     assert beamline.line[1] == "quad1"
+    assert beamline.line[1].name == "quad1"
+    assert beamline.line[1].element is None  # Not yet resolved
+    assert not beamline.line[1].is_resolved()
+
     # Third element should be a Drift object
     assert isinstance(beamline.line[2], pals.Drift)
     assert beamline.line[2].name == "drift2"
     assert beamline.line[2].length == 0.5
+
+    # Test that we can resolve the reference later
+    drift_element = pals.Drift(name="drift1", length=1.0)
+    beamline.line[0].element = drift_element
+    assert beamline.line[0].is_resolved()
+    assert beamline.line[0].element.name == "drift1"
+    assert beamline.line[0].element.length == 1.0
+
+
+def test_ElementReference_direct():
+    """Test ElementReference creation and behavior directly"""
+    # Test creation with positional argument
+    ref1 = pals.ElementReference("test_element")
+    assert ref1.name == "test_element"
+    assert str(ref1) == "test_element"
+    assert ref1 == "test_element"
+    assert not ref1.is_resolved()
+
+    # Test creation with keyword argument
+    ref2 = pals.ElementReference(name="another_element")
+    assert ref2.name == "another_element"
+    assert str(ref2) == "another_element"
+    assert ref2 == "another_element"
+
+    # Test hash (for use in sets/dicts)
+    ref_set = {ref1, ref2}
+    assert len(ref_set) == 2
+    assert ref1 in ref_set
+
+    # Test resolution
+    drift = pals.Drift(name="test_element", length=2.5)
+    ref1.element = drift
+    assert ref1.is_resolved()
+    assert ref1.element.length == 2.5
+
+    # Test repr
+    assert "test_element" in repr(ref1)
+    assert "resolved" in repr(ref1)
+    assert "unresolved" in repr(ref2)
+
+    # Test that element is a reference, not a copy
+    assert ref1.element is drift  # Same object identity
+    # Modify the original and verify the reference sees the change
+    drift.length = 3.0
+    assert ref1.element.length == 3.0  # Change is visible through reference