PALS Root Structure (#63)

Author: ax3l

examples/fodo.py

@@ -3,6 +3,7 @@
 from pals import Quadrupole
 from pals import BeamLine
 from pals import Lattice
+from pals import load, store
 
 
 def main():
@@ -51,23 +52,23 @@ def main():
 
     # Serialize to YAML
     yaml_file = "examples_fodo.pals.yaml"
-    lattice.to_file(yaml_file)
+    store(yaml_file, lattice)
 
     # Read YAML data from file
-    loaded_lattice = Lattice.from_file(yaml_file)
+    loaded_lattice = load(yaml_file)
 
     # Validate loaded data
-    assert lattice == loaded_lattice
+    assert lattice == loaded_lattice.facility[0]
 
     # Serialize to JSON
     json_file = "examples_fodo.pals.json"
-    lattice.to_file(json_file)
+    store(json_file, lattice)
 
     # Read JSON data from file
-    loaded_lattice = Lattice.from_file(json_file)
+    loaded_lattice = load(json_file)
 
     # Validate loaded data
-    assert lattice == loaded_lattice
+    assert lattice == loaded_lattice.facility[0]
 
 
 if __name__ == "__main__":

src/pals/PALS.py

@@ -0,0 +1,65 @@
+from pydantic import BaseModel
+
+from pydantic import model_validator
+from typing import List, Optional
+
+from .kinds import Lattice
+from .kinds.all_elements import get_all_elements_as_annotation
+from .functions import load_file_to_dict, store_dict_to_file
+
+
+Facility = List[get_all_elements_as_annotation()]
+
+
+class PALSroot(BaseModel):
+    """Represent the roo PALS structure"""
+
+    version: Optional[str] = None
+
+    facility: Facility
+
+    @model_validator(mode="before")
+    @classmethod
+    def unpack_json_structure(cls, data):
+        """Deserialize the JSON/YAML/...-like dict for facility elements"""
+        from pals.kinds.mixin.all_element_mixin import unpack_element_list_structure
+
+        return unpack_element_list_structure(data, "facility", "facility")
+
+    def model_dump(self, *args, **kwargs):
+        """Custom model dump for facility to handle element list formatting"""
+        from pals.kinds.mixin.all_element_mixin import dump_element_list
+
+        data = {}
+        data["PALS"] = {}
+        data["PALS"]["version"] = self.version
+        data["PALS"] = dump_element_list(self, "facility", *args, **kwargs)
+        return data
+
+    @staticmethod
+    def from_file(filename: str) -> "PALSroot":
+        """Load a facility from a text file"""
+        pals_dict = load_file_to_dict(filename)
+        return PALSroot(**pals_dict)
+
+    def to_file(self, filename: str):
+        """Save a facility to a text file"""
+        pals_dict = self.model_dump()
+        store_dict_to_file(filename, pals_dict)
+
+
+def load(filename: str) -> PALSroot:
+    """Load a facility from a text file"""
+    pals_dict = load_file_to_dict(filename)
+    return PALSroot(**pals_dict)
+
+
+def store(filename: str, pals_root: PALSroot | Facility | Lattice):
+    # wrap single elements in a list, facility in a PALSroot
+    if isinstance(pals_root, Lattice):
+        pals_root = PALSroot(version=None, facility=[pals_root])
+    elif isinstance(pals_root, list):
+        pals_root = PALSroot(version=None, facility=pals_root)
+
+    pals_dict = pals_root.model_dump()
+    store_dict_to_file(filename, pals_dict)

src/pals/__init__.py

@@ -6,3 +6,8 @@
 
 from .kinds import *  # noqa
 from .parameters import *  # noqa
+from .PALS import PALSroot, load, store  # noqa
+
+
+# Rebuild pydantic models that depend on other classes
+PALSroot.model_rebuild()

src/pals/kinds/mixin/all_element_mixin.py

@@ -110,5 +110,8 @@ def dump_element_list(self, field_name: str, *args, **kwargs) -> dict:
         elem_dict = elem.model_dump(**kwargs)
         new_list.append(elem_dict)
 
-    data[self.name][field_name] = new_list
+    if hasattr(self, "name"):  # all but PALSroot have a name
+        data[self.name][field_name] = new_list
+    else:
+        data[field_name] = new_list
     return data