Start draft of new Lattice class

- Copy implementation from BeamLine to Lattice
- Move methods from_file, to_file from BeamLine to Lattice
- Use new class in external examples test script

Author: EZoni

examples/test_external_examples.py

@@ -1,6 +1,6 @@
 import argparse
 
-from pals import BeamLine
+from pals import Lattice
 
 
 def main():
@@ -14,7 +14,7 @@ def main():
     args = parser.parse_args()
     example_file = args.path
     # Parse and validate YAML data from file
-    BeamLine.from_file(example_file)
+    Lattice.from_file(example_file)
 
 
 if __name__ == "__main__":

src/pals/kinds/BeamLine.py

@@ -3,7 +3,6 @@
 
 from .all_elements import get_all_elements_as_annotation
 from .mixin import BaseElement
-from ..functions import load_file_to_dict, store_dict_to_file
 
 
 class BeamLine(BaseElement):
@@ -26,14 +25,3 @@ def model_dump(self, *args, **kwargs):
         from pals.kinds.mixin.all_element_mixin import dump_element_list
 
         return dump_element_list(self, "line", *args, **kwargs)
-
-    @staticmethod
-    def from_file(filename: str) -> "BeamLine":
-        """Load a BeamLine from a text file"""
-        pals_dict = load_file_to_dict(filename)
-        return BeamLine(**pals_dict)
-
-    def to_file(self, filename: str):
-        """Save a BeamLine to a text file"""
-        pals_dict = self.model_dump()
-        store_dict_to_file(filename, pals_dict)

src/pals/kinds/Lattice.py

@@ -0,0 +1,39 @@
+from pydantic import model_validator
+from typing import List, Literal
+
+from .all_elements import get_all_elements_as_annotation
+from .mixin import BaseElement
+from ..functions import load_file_to_dict, store_dict_to_file
+
+
+class Lattice(BaseElement):
+    """A line of elements and/or other lines"""
+
+    kind: Literal["Lattice"] = "Lattice"
+
+    line: List[get_all_elements_as_annotation()]
+
+    @model_validator(mode="before")
+    @classmethod
+    def unpack_json_structure(cls, data):
+        """Deserialize the JSON/YAML/...-like dict for Lattice elements"""
+        from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
+
+        return unpack_element_list_structure(data, "line", "line")
+
+    def model_dump(self, *args, **kwargs):
+        """Custom model dump for Lattice to handle element list formatting"""
+        from pals.kinds.mixin.all_element_mixin import dump_element_list
+
+        return dump_element_list(self, "line", *args, **kwargs)
+
+    @staticmethod
+    def from_file(filename: str) -> "Lattice":
+        """Load a Lattice from a text file"""
+        pals_dict = load_file_to_dict(filename)
+        return Lattice(**pals_dict)
+
+    def to_file(self, filename: str):
+        """Save a Lattice to a text file"""
+        pals_dict = self.model_dump()
+        store_dict_to_file(filename, pals_dict)

src/pals/kinds/all_elements.py

@@ -43,6 +43,7 @@
 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 = (
+        "Lattice",  # Forward reference to handle circular import
         "BeamLine",  # Forward reference to handle circular import
         "UnionEle",  # Forward reference to handle circular import
         ACKicker,