Add PlaceholderName (#60)

Author: ax3l

src/pals/kinds/PlaceholderName.py

@@ -0,0 +1,84 @@
+from pydantic import BaseModel, Field, model_serializer
+from typing import Annotated
+
+from .mixin import BaseElement
+
+
+class PlaceholderName(BaseModel):
+    """Represents a reference to a named element.
+
+    This placeholder is replaced with a physically distinct
+    element during beamline expansion.
+
+    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 = PlaceholderName(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, PlaceholderName):
+            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 PlaceholderName."""
+        resolved = "resolved" if self.is_resolved() else "unresolved"
+        return f"PlaceholderName('{self.name}', {resolved})"

src/pals/kinds/__init__.py

@@ -26,6 +26,7 @@
 from .NullEle import NullEle  # noqa: F401
 from .Octupole import Octupole  # noqa: F401
 from .Patch import Patch  # noqa: F401
+from .PlaceholderName import PlaceholderName  # noqa: F401
 from .Quadrupole import Quadrupole  # noqa: F401
 from .RBend import RBend  # noqa: F401
 from .RFCavity import RFCavity  # 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 .PlaceholderName import PlaceholderName
 from .Feedback import Feedback
 from .Fiducial import Fiducial
 from .FloorShift import FloorShift
@@ -83,6 +83,13 @@ 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."""
-    types = get_all_element_types(extra_types)
-    return Annotated[Union[types], Field(discriminator="kind")]
+    """Return the Union type of all allowed elements with their kind as the discriminator field.
+
+    Note: PlaceholderName is included to support string references to named elements.
+    Since PlaceholderName 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) + (PlaceholderName,)
+    # We can't use discriminator with PlaceholderName 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 ..PlaceholderName import PlaceholderName
 
 
 def unpack_element_list_structure(
@@ -44,7 +45,14 @@ 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):
-            raise RuntimeError("Reference/alias elements not yet implemented")
+            # Wrap the string in a Placeholder name object
+            new_list.append(PlaceholderName(item))
+            continue
+        # An element can be a PlaceholderName instance directly
+        elif isinstance(item, PlaceholderName):
+            # Keep the PlaceholderName as-is
+            new_list.append(item)
+            continue
         # An element can be a dict
         elif isinstance(item, dict):
             if not (len(item) == 1):
@@ -69,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, PlaceholderName, 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

@@ -537,3 +537,90 @@ def test_Lattice():
     assert lattice.branches[1].name == "line2"
     assert lattice.branches[0].line == [element1, element2, element3]
     assert lattice.branches[1].line == [element3]
+
+
+def test_BeamLine_with_string_references():
+    """Test BeamLine with string references to elements defined elsewhere"""
+    import yaml
+
+    # Test YAML deserialization with string references
+    yaml_data = """
+    fodo_cell:
+        kind: BeamLine
+        line:
+          - drift1
+          - quad1
+          - drift2:
+              kind: Drift
+              length: 0.5
+    """
+
+    data = yaml.safe_load(yaml_data)
+    beamline = pals.BeamLine(**data)
+
+    assert beamline.name == "fodo_cell"
+    assert len(beamline.line) == 3
+
+    # First element should be a PlaceholderName that behaves like the string "drift1"
+    assert isinstance(beamline.line[0], pals.PlaceholderName)
+    assert beamline.line[0] == "drift1"
+    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 a PlaceholderName that behaves like the string "quad1"
+    assert isinstance(beamline.line[1], pals.PlaceholderName)
+    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_PlaceholderName_direct():
+    """Test PlaceholderName creation and behavior directly"""
+    # Test creation with positional argument
+    ref1 = pals.PlaceholderName("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.PlaceholderName(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