Alternative implementation of name based listing

Author: EZoni

examples/fodo.py

@@ -43,16 +43,16 @@ def main():
     )
     # Create line with all elements
     line = Line(
-        line=[
-            drift1,
-            quad1,
-            drift2,
-            quad2,
-            drift3,
-        ]
+        line={
+            drift1.name: drift1,
+            quad1.name: quad1,
+            drift2.name: drift2,
+            quad2.name: quad2,
+            drift3.name: drift3,
+        }
     )
     # Serialize to YAML
-    yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
+    yaml_data = yaml.dump(line.model_dump(), default_flow_style=False, sort_keys=False)
     print("Dumping YAML data...")
     print(f"{yaml_data}")
     # Write YAML data to file
@@ -67,7 +67,7 @@ def main():
     # Validate loaded data
     assert line == loaded_line
     # Serialize to JSON
-    json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2)
+    json_data = json.dumps(line.model_dump(), indent=2, sort_keys=False)
     print("Dumping JSON data...")
     print(f"{json_data}")
     # Write JSON data to file

schema/Line.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel, ConfigDict, Field
-from typing import Annotated, List, Literal, Union
+from typing import Annotated, Literal, OrderedDict, Union
 
 from schema.BaseElement import BaseElement
 from schema.ThickElement import ThickElement
@@ -16,7 +16,8 @@ class Line(BaseModel):
 
     kind: Literal["Line"] = "Line"
 
-    line: List[
+    line: OrderedDict[
+        str,
         Annotated[
             Union[
                 BaseElement,
@@ -26,7 +27,7 @@ class Line(BaseModel):
                 "Line",
             ],
             Field(discriminator="kind"),
-        ]
+        ],
     ]
 
 

tests/test_schema.py

@@ -104,18 +104,22 @@ def test_QuadrupoleElement():
 def test_Line():
     # Create first line with one base element
     element1 = BaseElement(name="element1")
-    line1 = Line(line=[element1])
-    assert line1.line == [element1]
+    line1 = Line(line={element1.name: element1})
+    assert line1.line == {element1.name: element1}
     # Extend first line with one thick element
     element2 = ThickElement(name="element2", length=2.0)
-    line1.line.extend([element2])
-    assert line1.line == [element1, element2]
+    line1.line.update({element2.name: element2})
+    assert line1.line == {element1.name: element1, element2.name: element2}
     # Create second line with one drift element
     element3 = DriftElement(name="element3", length=3.0)
-    line2 = Line(line=[element3])
+    line2 = Line(line={element3.name: element3})
     # Extend first line with second line
-    line1.line.extend(line2.line)
-    assert line1.line == [element1, element2, element3]
+    line1.line.update(line2.line)
+    assert line1.line == {
+        element1.name: element1,
+        element2.name: element2,
+        element3.name: element3,
+    }
 
 
 def test_yaml():
@@ -124,7 +128,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(line={element1.name: element1, element2.name: element2})
     # Serialize the Line object to YAML
     yaml_data = yaml.dump(line.model_dump(), default_flow_style=False)
     print(f"\n{yaml_data}")
@@ -149,7 +153,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(line={element1.name: element1, element2.name: element2})
     # Serialize the Line object to JSON
     json_data = json.dumps(line.model_dump(), sort_keys=True, indent=2)
     print(f"\n{json_data}")