Fix line validation

Author: EZoni

examples/fodo.py

@@ -64,7 +64,7 @@ def main():
     with open(yaml_file, "r") as file:
         yaml_data = yaml.safe_load(file)
     # Parse YAML data
-    loaded_line = Line(**yaml_data)
+    loaded_line = Line(**yaml_data[0])
     # Validate loaded data
     assert line == loaded_line
     # Serialize to JSON
@@ -79,7 +79,7 @@ def main():
     with open(json_file, "r") as file:
         json_data = json.loads(file.read())
     # Parse JSON data
-    loaded_line = Line(**json_data)
+    loaded_line = Line(**json_data[0])
     # Validate loaded data
     assert line == loaded_line
 

schema/Line.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import ConfigDict, Field, model_validator
 from typing import Annotated, List, Literal, Union
 
 from schema.BaseElement import BaseElement
@@ -29,47 +29,70 @@ class Line(BaseElement):
         ]
     ]
 
-    @field_validator("line", mode="before")
+    # @field_validator("line", mode="before")
+    # @classmethod
+    # def parse_list_of_dicts(cls, value):
+    #    """This method inserts the key of the one-key dictionary into
+    #    the name attribute of the elements"""
+    #    if not isinstance(value, list):
+    #        raise TypeError("line must be a list")
+
+    #    if value and isinstance(value[0], BaseModel):
+    #        # Already a list of models; nothing to do
+    #        return value
+
+    #    # we expect a list of dicts or strings
+    #    elements = []
+    #    for item_dict in value:
+    #        # an element is either a reference string to another element or a dict
+    #        if isinstance(item_dict, str):
+    #            raise RuntimeError("Reference/alias elements not yet implemented")
+
+    #        elif isinstance(item_dict, dict):
+    #            if not (isinstance(item_dict, dict) and len(item_dict) == 1):
+    #                raise ValueError(
+    #                    f"Each line element must be a dict with exactly one key, the name of the element, but we got: {item_dict!r}"
+    #                )
+    #            [(name, fields)] = item_dict.items()
+
+    #            if not isinstance(fields, dict):
+    #                raise ValueError(
+    #                    f"Value for element key '{name}' must be a dict (got {fields!r})"
+    #                )
+
+    #        # Insert the name into the fields dict
+    #        fields["name"] = name
+    #        elements.append(fields)
+    #    return elements
+
+    @model_validator(mode="before")
     @classmethod
-    def parse_list_of_dicts(cls, value):
-        """This method inserts the key of the one-key dictionary into
-        the name attribute of the elements"""
-        if not isinstance(value, list):
-            raise TypeError("line must be a list")
-
-        if value and isinstance(value[0], BaseModel):
-            # Already a list of models; nothing to do
-            return value
-
-        # we expect a list of dicts or strings
-        elements = []
-        for item_dict in value:
-            # an element is either a reference string to another element or a dict
-            if isinstance(item_dict, str):
-                raise RuntimeError("Reference/alias elements not yet implemented")
-
-            elif isinstance(item_dict, dict):
-                if not (isinstance(item_dict, dict) and len(item_dict) == 1):
-                    raise ValueError(
-                        f"Each line element must be a dict with exactly one key, the name of the element, but we got: {item_dict!r}"
-                    )
-                [(name, fields)] = item_dict.items()
-
-                if not isinstance(fields, dict):
-                    raise ValueError(
-                        f"Value for element key '{name}' must be a dict (got {fields!r})"
-                    )
-
-            # Insert the name into the fields dict
-            fields["name"] = name
-            elements.append(fields)
-        return elements
+    def unpack_yaml_structure(cls, data):
+        # Handle the top-level dict
+        if isinstance(data, dict) and len(data) == 1:
+            name, value = list(data.items())[0]
+            value = dict(value)
+            value["name"] = name
+            data = value
+        # Now handle the 'line' field if present: unpack each element's name
+        if "line" in data and isinstance(data["line"], list):
+            new_line = []
+            for item in data["line"]:
+                if isinstance(item, dict) and len(item) == 1:
+                    elem_name, elem_fields = list(item.items())[0]
+                    elem_fields = dict(elem_fields)
+                    elem_fields["name"] = elem_name
+                    new_line.append(elem_fields)
+                else:
+                    new_line.append(item)
+            data["line"] = new_line
+        return data
 
     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 one-element list of the form
         # [{key: value}], where 'key' is the name of the line and 'value' is a
-        # dictionary with all other properties
+        # dict with all other properties
         data = super().model_dump(*args, **kwargs)
         # Reformat 'line' field as list of single-key dicts
         new_line = []

tests/test_schema.py

@@ -104,15 +104,15 @@ def test_QuadrupoleElement():
 def test_Line():
     # Create first line with one base element
     element1 = BaseElement(name="element1")
-    line1 = Line(line=[element1])
+    line1 = Line(name="line1", line=[element1])
     assert line1.line == [element1]
     # Extend first line with one thick element
     element2 = ThickElement(name="element2", length=2.0)
     line1.line.extend([element2])
     assert line1.line == [element1, element2]
     # Create second line with one drift element
     element3 = DriftElement(name="element3", length=3.0)
-    line2 = Line(line=[element3])
+    line2 = Line(name="line2", line=[element3])
     # Extend first line with second line
     line1.line.extend(line2.line)
     assert line1.line == [element1, element2, element3]
@@ -124,7 +124,7 @@ def test_yaml():
     # Create one thick element
     element2 = ThickElement(name="element2", length=2.0)
     # Create line with both elements
-    line = Line(line=[element1, element2])
+    line = Line(name="line", line=[element1, element2])
     # Serialize the Line object to YAML
     yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
     print(f"\n{yaml_data}")
@@ -136,7 +136,7 @@ def test_yaml():
     with open(test_file, "r") as file:
         yaml_data = yaml.safe_load(file)
     # Parse the YAML data back into a Line object
-    loaded_line = Line(**yaml_data)
+    loaded_line = Line(**yaml_data[0])
     # Remove the test file
     os.remove(test_file)
     # Validate loaded Line object
@@ -149,7 +149,7 @@ def test_json():
     # Create one thick element
     element2 = ThickElement(name="element2", length=2.0)
     # Create line with both elements
-    line = Line(line=[element1, element2])
+    line = Line(name="line", line=[element1, element2])
     # Serialize the Line object to JSON
     json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2)
     print(f"\n{json_data}")
@@ -161,7 +161,7 @@ def test_json():
     with open(test_file, "r") as file:
         json_data = json.loads(file.read())
     # Parse the JSON data back into a Line object
-    loaded_line = Line(**json_data)
+    loaded_line = Line(**json_data[0])
     # Remove the test file
     os.remove(test_file)
     # Validate loaded Line object