More Parameter Validations

And more limits on ints and string literal values.

Author: ax3l

src/pals/parameters/ApertureParameters.py

@@ -1,15 +1,25 @@
-from typing import List
-from pydantic import BaseModel
+from typing import Literal
+from pydantic import BaseModel, Field
 
 
 class ApertureParameters(BaseModel):
     """Aperture parameters"""
 
-    x_limits: List[float] = [float("nan"), float("nan")]
-    y_limits: List[float] = [float("nan"), float("nan")]
-    shape: str = "RECTANGULAR"
-    location: str = "ENTRANCE_END"
+    x_limits: list[float | None, float | None] = Field(
+        default=[None, None],
+        validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]),
+    )
+    y_limits: list[float | None, float | None] = Field(
+        default=[None, None],
+        validate_args=lambda x: (x[0] is None or x[1] is None or x[0] < x[1]),
+    )
+    shape: Literal["RECTANGULAR", "ELLIPTICAL", "VERTICES", "CUSTOM_SHAPE"] = (
+        "RECTANGULAR"
+    )
+    location: Literal[
+        "ENTRANCE_END", "CENTER", "EXIT_END", "BOTH_ENDS", "NOWHERE", "EVERYWHERE"
+    ] = "ENTRANCE_END"
     material: str = ""
-    thickness: float = 0.0
+    thickness: float = Field(default=0.0, ge=0.0)
     aperture_shifts_with_body: bool = False
     aperture_active: bool = True

src/pals/parameters/ElectricMultipoleParameters.py

@@ -4,7 +4,7 @@
 class ElectricMultipoleParameters(BaseModel):
     """Electric multipole parameters"""
 
-    # Allow arbitrary fields
+    # Allow arbitrary fields (TODO: remove this)
     model_config = ConfigDict(extra="allow")
 
     # TODO: add ElectricMultipoleParameters in a follow-up RP

src/pals/parameters/ForkParameters.py

@@ -1,3 +1,4 @@
+from typing import Literal
 from pydantic import BaseModel
 
 
@@ -6,5 +7,5 @@ class ForkParameters(BaseModel):
 
     to_line: str = ""
     to_ele: str = ""
-    direction: str = "FORWARDS"  # "FORWARDS" or "BACKWARDS"
+    direction: Literal["FORWARDS", "BACKWARDS"] = "FORWARDS"
     propagate_reference: bool = True

src/pals/parameters/PatchParameters.py

@@ -1,3 +1,4 @@
+from typing import Literal
 from pydantic import BaseModel
 
 
@@ -11,5 +12,5 @@ class PatchParameters(BaseModel):
     y_rot: float = 0.0
     z_rot: float = 0.0
     flexible: bool = False
-    ref_coords: str = "exit_end"  # "entrance_end" or "exit_end"
+    ref_coords: Literal["entrance_end", "exit_end"] = "exit_end"
     user_sets_length: bool = False

src/pals/parameters/RFParameters.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class RFParameters(BaseModel):
@@ -11,4 +11,4 @@ class RFParameters(BaseModel):
     phase: float = 0.0  # [unitless] RF phase in 0 to 2*pi
     multipass_phase: float = 0.0  # [unitless] RF Phase added to multipass elements
     cavity_type: str = "STANDING_WAVE"  # [string] Cavity type
-    n_cell: int = 1  # [unitless] Number of cavity cells
+    n_cell: int = Field(default=1, gt=0)  # [unitless] Number of cavity cells

src/pals/parameters/ReferenceParameters.py

@@ -1,3 +1,4 @@
+from typing import Literal
 from pydantic import BaseModel
 
 
@@ -9,3 +10,6 @@ class ReferenceParameters(BaseModel):
     E_tot_ref: float = 0.0  # [eV] Reference total energy
     time_ref: float = 0.0  # [s] Reference time
     location: str = ""  # Where reference parameters are evaluated
+    location: Literal[
+        "UPSTREAM_END", "DOWNSTREAM_END"
+    ]  # TODO: undefined default in PALS?

tests/test_elements.py

@@ -1,3 +1,4 @@
+import pytest
 from pydantic import ValidationError
 
 import pals
@@ -23,13 +24,8 @@ def test_ThickElement():
     # Try to assign negative length and
     # detect validation error without breaking pytest
     element_length = -1.0
-    passed = True
-    try:
+    with pytest.raises(ValidationError):
         element.length = element_length
-    except ValidationError as e:
-        print(e)
-        passed = False
-    assert not passed
 
 
 def test_Drift():
@@ -45,13 +41,8 @@ def test_Drift():
     # Try to assign negative length and
     # detect validation error without breaking pytest
     element_length = -1.0
-    passed = True
-    try:
+    with pytest.raises(ValidationError):
         element.length = element_length
-    except ValidationError as e:
-        print(e)
-        passed = False
-    assert not passed
 
 
 def test_Quadrupole():

tests/test_parameters.py

@@ -1,10 +1,11 @@
+import pytest
+from pydantic import ValidationError
+
 from pals import (
     ApertureParameters,
     BeamBeamParameters,
     BendParameters,
     BodyShiftParameters,
-    ElectricMultipoleParameters,
-    # FloorParameters,  # not yet tested
     FloorShiftParameters,
     ForkParameters,
     MagneticMultipoleParameters,
@@ -32,15 +33,23 @@ def test_ParameterClasses():
     meta = MetaParameters(alias="test", description="test element")
     assert meta.alias == "test"
 
-    # Test ElectricMultipoleParameters
-    emp = ElectricMultipoleParameters(En1=1.0, Es1=0.5)
-    assert emp.En1 == 1.0
+    # Test ElectricMultipoleParameters (TODO)
+    # emp = ElectricMultipoleParameters(En1=1.0, Es1=0.5)
+    # assert emp.En1 == 1.0
 
     # Test MagneticMultipoleParameters
     mmp = MagneticMultipoleParameters(Bn1=1.0, Bs1=0.5)
     assert mmp.Bn1 == 1.0
     assert mmp.Bs1 == 0.5
 
+    #  catch typos
+    with pytest.raises(ValidationError):
+        _ = MagneticMultipoleParameters(Bm1=1.0, Bs1=0.5)
+    with pytest.raises(ValidationError):
+        _ = MagneticMultipoleParameters(Bn1=1.0, Bv1=0.5)
+    with pytest.raises(ValidationError):
+        _ = MagneticMultipoleParameters(Bn01=1.0, Bs01=0.5)
+
     # Test SolenoidParameters
     sol = SolenoidParameters(Ksol=0.1, Bsol=0.2)
     assert sol.Ksol == 0.1
@@ -49,6 +58,11 @@ def test_ParameterClasses():
     rf = RFParameters(frequency=1e9, voltage=1e6)
     assert rf.frequency == 1e9
 
+    with pytest.raises(ValidationError):
+        _ = RFParameters(frequency=1e9, voltage=1e6, n_cell=0)
+    with pytest.raises(ValidationError):
+        _ = RFParameters(frequency=1e9, voltage=1e6, n_cell=-1)
+
     # Test BendParameters
     bend = BendParameters(rho_ref=1.0, bend_field_ref=2.0)
     assert bend.rho_ref == 1.0