From e6ba8401a0de947838c6efd8843c3a8cd3e02956 Mon Sep 17 00:00:00 2001 From: Edoardo Zoni Date: Wed, 15 Oct 2025 11:05:08 -0700 Subject: [PATCH] Skip 'Element' suffix for common elements --- examples/fodo.py | 14 +++++++------- src/pals_schema/BeamLine.py | 8 ++++---- src/pals_schema/{DriftElement.py => Drift.py} | 2 +- .../{QuadrupoleElement.py => Quadrupole.py} | 2 +- tests/test_schema.py | 14 +++++++------- 5 files changed, 20 insertions(+), 20 deletions(-) rename src/pals_schema/{DriftElement.py => Drift.py} (82%) rename src/pals_schema/{QuadrupoleElement.py => Quadrupole.py} (89%) diff --git a/examples/fodo.py b/examples/fodo.py index 223f084..3384be6 100644 --- a/examples/fodo.py +++ b/examples/fodo.py @@ -7,35 +7,35 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../src/"))) from pals_schema.MagneticMultipoleParameters import MagneticMultipoleParameters -from pals_schema.DriftElement import DriftElement -from pals_schema.QuadrupoleElement import QuadrupoleElement +from pals_schema.Drift import Drift +from pals_schema.Quadrupole import Quadrupole from pals_schema.BeamLine import BeamLine def main(): - drift1 = DriftElement( + drift1 = Drift( name="drift1", length=0.25, ) - quad1 = QuadrupoleElement( + quad1 = Quadrupole( name="quad1", length=1.0, MagneticMultipoleP=MagneticMultipoleParameters( Bn1=1.0, ), ) - drift2 = DriftElement( + drift2 = Drift( name="drift2", length=0.5, ) - quad2 = QuadrupoleElement( + quad2 = Quadrupole( name="quad2", length=1.0, MagneticMultipoleP=MagneticMultipoleParameters( Bn1=-1.0, ), ) - drift3 = DriftElement( + drift3 = Drift( name="drift3", length=0.5, ) diff --git a/src/pals_schema/BeamLine.py b/src/pals_schema/BeamLine.py index 29f035a..f2dd400 100644 --- a/src/pals_schema/BeamLine.py +++ b/src/pals_schema/BeamLine.py @@ -3,8 +3,8 @@ from pals_schema.BaseElement import BaseElement from pals_schema.ThickElement import ThickElement -from pals_schema.DriftElement import DriftElement -from pals_schema.QuadrupoleElement import QuadrupoleElement +from pals_schema.Drift import Drift +from pals_schema.Quadrupole import Quadrupole class BeamLine(BaseElement): @@ -21,8 +21,8 @@ class BeamLine(BaseElement): Union[ BaseElement, ThickElement, - DriftElement, - QuadrupoleElement, + Drift, + Quadrupole, "BeamLine", ], Field(discriminator="kind"), diff --git a/src/pals_schema/DriftElement.py b/src/pals_schema/Drift.py similarity index 82% rename from src/pals_schema/DriftElement.py rename to src/pals_schema/Drift.py index ce1bc01..18bcfa9 100644 --- a/src/pals_schema/DriftElement.py +++ b/src/pals_schema/Drift.py @@ -3,7 +3,7 @@ from .ThickElement import ThickElement -class DriftElement(ThickElement): +class Drift(ThickElement): """A field free region""" # Discriminator field diff --git a/src/pals_schema/QuadrupoleElement.py b/src/pals_schema/Quadrupole.py similarity index 89% rename from src/pals_schema/QuadrupoleElement.py rename to src/pals_schema/Quadrupole.py index f778a3d..127b63e 100644 --- a/src/pals_schema/QuadrupoleElement.py +++ b/src/pals_schema/Quadrupole.py @@ -4,7 +4,7 @@ from .MagneticMultipoleParameters import MagneticMultipoleParameters -class QuadrupoleElement(ThickElement): +class Quadrupole(ThickElement): """A quadrupole element""" # Discriminator field diff --git a/tests/test_schema.py b/tests/test_schema.py index 393bf5c..3b8fd90 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -11,8 +11,8 @@ from pals_schema.MagneticMultipoleParameters import MagneticMultipoleParameters from pals_schema.BaseElement import BaseElement from pals_schema.ThickElement import ThickElement -from pals_schema.DriftElement import DriftElement -from pals_schema.QuadrupoleElement import QuadrupoleElement +from pals_schema.Drift import Drift +from pals_schema.Quadrupole import Quadrupole from pals_schema.BeamLine import BeamLine @@ -45,11 +45,11 @@ def test_ThickElement(): assert not passed -def test_DriftElement(): +def test_Drift(): # Create one drift element with custom name and length element_name = "drift_element" element_length = 1.0 - element = DriftElement( + element = Drift( name=element_name, length=element_length, ) @@ -67,7 +67,7 @@ def test_DriftElement(): assert not passed -def test_QuadrupoleElement(): +def test_Quadrupole(): # Create one drift element with custom name and length element_name = "quadrupole_element" element_length = 1.0 @@ -85,7 +85,7 @@ def test_QuadrupoleElement(): Bs2=element_magnetic_multipole_Bs2, tilt2=element_magnetic_multipole_tilt2, ) - element = QuadrupoleElement( + element = Quadrupole( name=element_name, length=element_length, MagneticMultipoleP=element_magnetic_multipole, @@ -113,7 +113,7 @@ def test_BeamLine(): line1.line.extend([element2]) assert line1.line == [element1, element2] # Create second line with one drift element - element3 = DriftElement(name="element3", length=3.0) + element3 = Drift(name="element3", length=3.0) line2 = BeamLine(name="line2", line=[element3]) # Extend first line with second line line1.line.extend(line2.line)