diff --git a/examples/fodo.py b/examples/fodo.py index c2e6ed5..d67eaaa 100644 --- a/examples/fodo.py +++ b/examples/fodo.py @@ -22,7 +22,7 @@ def main(): quad1 = QuadrupoleElement( name="quad1", length=1.0, - MagneticMultipoleP=MagneticMultipoleParameters( + magnetic_multipole_parameters=MagneticMultipoleParameters( Bn1=1.0, ), ) @@ -33,7 +33,7 @@ def main(): quad2 = QuadrupoleElement( name="quad2", length=1.0, - MagneticMultipoleP=MagneticMultipoleParameters( + magnetic_multipole_parameters=MagneticMultipoleParameters( Bn1=-1.0, ), ) diff --git a/schema/BaseElement.py b/schema/BaseElement.py index 825635e..3608e46 100644 --- a/schema/BaseElement.py +++ b/schema/BaseElement.py @@ -6,7 +6,7 @@ class BaseElement(BaseModel): """A custom base element defining common properties""" # Discriminator field - kind: Literal["BaseElement"] = "BaseElement" + Type: Literal["BaseElement"] = "BaseElement" # Validate every time a new value is assigned to an attribute, # not only when an instance of BaseElement is created diff --git a/schema/DriftElement.py b/schema/DriftElement.py index ce1bc01..5634f63 100644 --- a/schema/DriftElement.py +++ b/schema/DriftElement.py @@ -7,4 +7,4 @@ class DriftElement(ThickElement): """A field free region""" # Discriminator field - kind: Literal["Drift"] = "Drift" + Type: Literal["Drift"] = "Drift" diff --git a/schema/Line.py b/schema/Line.py index e1ad070..0ebc2e8 100644 --- a/schema/Line.py +++ b/schema/Line.py @@ -14,7 +14,7 @@ class Line(BaseModel): # not only when an instance of Line is created model_config = ConfigDict(validate_assignment=True) - kind: Literal["Line"] = "Line" + Type: Literal["Line"] = "Line" line: List[ Annotated[ @@ -25,7 +25,7 @@ class Line(BaseModel): QuadrupoleElement, "Line", ], - Field(discriminator="kind"), + Field(discriminator="Type"), ] ] diff --git a/schema/QuadrupoleElement.py b/schema/QuadrupoleElement.py index f778a3d..6dda0e8 100644 --- a/schema/QuadrupoleElement.py +++ b/schema/QuadrupoleElement.py @@ -8,7 +8,7 @@ class QuadrupoleElement(ThickElement): """A quadrupole element""" # Discriminator field - kind: Literal["Quadrupole"] = "Quadrupole" + Type: Literal["Quadrupole"] = "Quadrupole" # Magnetic multipole parameters - MagneticMultipoleP: MagneticMultipoleParameters + magnetic_multipole_parameters: MagneticMultipoleParameters diff --git a/schema/ThickElement.py b/schema/ThickElement.py index 9461041..4e5b584 100644 --- a/schema/ThickElement.py +++ b/schema/ThickElement.py @@ -8,7 +8,7 @@ class ThickElement(BaseElement): """A thick base element with finite segment length""" # Discriminator field - kind: Literal["ThickElement"] = "ThickElement" + Type: Literal["ThickElement"] = "ThickElement" # Segment length in meters (m) length: Annotated[float, Gt(0)] diff --git a/tests/test_schema.py b/tests/test_schema.py index 3abdc8a..7e5bfb7 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -75,7 +75,7 @@ def test_QuadrupoleElement(): element_magnetic_multipole_Bs2 = 2.2 element_magnetic_multipole_tilt1 = 3.1 element_magnetic_multipole_tilt2 = 3.2 - element_magnetic_multipole = MagneticMultipoleParameters( + element_magnetic_multipole_parameters = MagneticMultipoleParameters( Bn1=element_magnetic_multipole_Bn1, Bs1=element_magnetic_multipole_Bs1, tilt1=element_magnetic_multipole_tilt1, @@ -86,16 +86,20 @@ def test_QuadrupoleElement(): element = QuadrupoleElement( name=element_name, length=element_length, - MagneticMultipoleP=element_magnetic_multipole, + magnetic_multipole_parameters=element_magnetic_multipole_parameters, ) assert element.name == element_name assert element.length == element_length - assert element.MagneticMultipoleP.Bn1 == element_magnetic_multipole_Bn1 - assert element.MagneticMultipoleP.Bs1 == element_magnetic_multipole_Bs1 - assert element.MagneticMultipoleP.tilt1 == element_magnetic_multipole_tilt1 - assert element.MagneticMultipoleP.Bn2 == element_magnetic_multipole_Bn2 - assert element.MagneticMultipoleP.Bs2 == element_magnetic_multipole_Bs2 - assert element.MagneticMultipoleP.tilt2 == element_magnetic_multipole_tilt2 + assert element.magnetic_multipole_parameters.Bn1 == element_magnetic_multipole_Bn1 + assert element.magnetic_multipole_parameters.Bs1 == element_magnetic_multipole_Bs1 + assert ( + element.magnetic_multipole_parameters.tilt1 == element_magnetic_multipole_tilt1 + ) + assert element.magnetic_multipole_parameters.Bn2 == element_magnetic_multipole_Bn2 + assert element.magnetic_multipole_parameters.Bs2 == element_magnetic_multipole_Bs2 + assert ( + element.magnetic_multipole_parameters.tilt2 == element_magnetic_multipole_tilt2 + ) # Serialize the Line object to YAML yaml_data = yaml.dump(element.model_dump(), default_flow_style=False) print(f"\n{yaml_data}")