diff --git a/pyaml/arrays/array.py b/pyaml/arrays/array.py index e3b59b73..a1e4e427 100644 --- a/pyaml/arrays/array.py +++ b/pyaml/arrays/array.py @@ -2,18 +2,17 @@ Array configuration """ -from pydantic import BaseModel, ConfigDict - from pyaml.common.exception import PyAMLException +from ..common.element import __pyaml_repr__ from ..common.element_holder import ElementHolder +from ..validation import DynamicValidation, register_schema -class ArrayConfigModel(BaseModel): +@register_schema +class ArrayConfig(DynamicValidation): """ - Base class for configuration of array of :py:class:`~pyaml.arrays.element.Element`, - :py:class:`~pyaml.arrays.bpm.BPM`, :py:class:`~pyaml.arrays.magnet.Magnet` or - :py:class:`~pyaml.arrays.cfm_magnet.CombinedFunctionMagnet`. + Base class for configuration of arrays (families). Parameters ---------- @@ -23,19 +22,17 @@ class ArrayConfigModel(BaseModel): List of pyaml element names """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - name: str - elements: list[str] - + def __init__(self, name: str, elements: list[str]): + self._name = name + self._elements = elements -class ArrayConfig(object): - """ - Base class that implements configuration for access to arrays (families) - """ + @property + def name(self): + return self._name - def __init__(self, cfg: ArrayConfigModel): - self._cfg = cfg + @property + def elements(self): + return self._elements def fill_array(self, holder: ElementHolder): """ @@ -57,6 +54,4 @@ def fill_array(self, holder: ElementHolder): raise PyAMLException("Array.fill_array() is not subclassed") def __repr__(self): - # ArrayConfigModel is a super class - # ConfigModel is expected from sub classes - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/pyaml/arrays/bpm.py b/pyaml/arrays/bpm.py index 0883a378..73c0ef80 100644 --- a/pyaml/arrays/bpm.py +++ b/pyaml/arrays/bpm.py @@ -1,15 +1,13 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from ..validation import DynamicValidation, register_schema +from .array import ArrayConfig # Define the main class name for this module PYAMLCLASS = "BPM" -class ConfigModel(ArrayConfigModel): - """Configuration model for :py:class:`.BPMArray`.""" - - -class BPM(ArrayConfig): +@register_schema +class BPM(ArrayConfig, DynamicValidation): """ :py:class:`.BPMArray` configuration. @@ -40,8 +38,8 @@ class BPM(ArrayConfig): """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__(self, name: str, elements: list[str]): + super().__init__(name, elements) def fill_array(self, holder: ElementHolder): """ @@ -66,4 +64,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate the :py:class:`.BPMArray` with. """ - holder.fill_bpm_array(self._cfg.name, self._cfg.elements) + holder.fill_bpm_array(self._name, self._elements) diff --git a/pyaml/arrays/cfm_magnet.py b/pyaml/arrays/cfm_magnet.py index f69ce658..a2cb38c3 100644 --- a/pyaml/arrays/cfm_magnet.py +++ b/pyaml/arrays/cfm_magnet.py @@ -1,17 +1,13 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from ..validation import DynamicValidation, register_schema +from .array import ArrayConfig # Define the main class name for this module PYAMLCLASS = "CombinedFunctionMagnet" -class ConfigModel(ArrayConfigModel): - """Configuration model for Combined Function Magnet array.""" - - ... - - -class CombinedFunctionMagnet(ArrayConfig): +@register_schema +class CombinedFunctionMagnet(ArrayConfig, DynamicValidation): """ Combined function magnet array confirguration @@ -28,8 +24,8 @@ class CombinedFunctionMagnet(ArrayConfig): ) """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__(self, name: str, elements: list[str]): + super().__init__(name, elements) def fill_array(self, holder: ElementHolder): """ @@ -40,4 +36,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with combined function magnet array """ - holder.fill_cfm_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_cfm_magnet_array(self._name, self._elements) diff --git a/pyaml/arrays/element.py b/pyaml/arrays/element.py index 6a40d027..2233ccc0 100644 --- a/pyaml/arrays/element.py +++ b/pyaml/arrays/element.py @@ -1,15 +1,13 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from ..validation import DynamicValidation, register_schema +from .array import ArrayConfig # Define the main class name for this module PYAMLCLASS = "Element" -class ConfigModel(ArrayConfigModel): - """Configuration model for :py:class:`.ElementArray`.""" - - -class Element(ArrayConfig): +@register_schema +class Element(ArrayConfig, DynamicValidation): """ :py:class:`.ElementArray` configuration. @@ -29,8 +27,8 @@ class Element(ArrayConfig): """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__(self, name: str, elements: list[str]): + super().__init__(name, elements) def fill_array(self, holder: ElementHolder): """ @@ -54,4 +52,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with element array """ - holder.fill_element_array(self._cfg.name, self._cfg.elements) + holder.fill_element_array(self._name, self._elements) diff --git a/pyaml/arrays/magnet.py b/pyaml/arrays/magnet.py index d65282cd..3c2d4028 100644 --- a/pyaml/arrays/magnet.py +++ b/pyaml/arrays/magnet.py @@ -1,17 +1,13 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from ..validation import DynamicValidation, register_schema +from .array import ArrayConfig # Define the main class name for this module PYAMLCLASS = "Magnet" -class ConfigModel(ArrayConfigModel): - """Configuration model for Magnet array.""" - - ... - - -class Magnet(ArrayConfig): +@register_schema +class Magnet(ArrayConfig, DynamicValidation): """ Magnet array confirguration @@ -27,8 +23,8 @@ class Magnet(ArrayConfig): ) """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__(self, name: str, elements: list[str]): + super().__init__(name, elements) def fill_array(self, holder: ElementHolder): """ @@ -39,4 +35,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with magnet array """ - holder.fill_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_magnet_array(self._name, self._elements) diff --git a/pyaml/arrays/serialized_magnet.py b/pyaml/arrays/serialized_magnet.py index fd412403..e055b3f1 100644 --- a/pyaml/arrays/serialized_magnet.py +++ b/pyaml/arrays/serialized_magnet.py @@ -1,17 +1,13 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from ..validation import DynamicValidation, register_schema +from .array import ArrayConfig # Define the main class name for this module PYAMLCLASS = "SerializedMagnets" -class ConfigModel(ArrayConfigModel): - """Configuration model for Serialized Magnets array.""" - - ... - - -class SerializedMagnets(ArrayConfig): +@register_schema +class SerializedMagnets(ArrayConfig, DynamicValidation): """ Serialized magnets array configuration @@ -28,8 +24,8 @@ class SerializedMagnets(ArrayConfig): ) """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__(self, name: str, elements: list[str]): + super().__init__(name, elements) def fill_array(self, holder: ElementHolder): """ @@ -40,4 +36,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with serialized magnet array """ - holder.fill_serialized_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_serialized_magnet_array(self._name, self._elements) diff --git a/pyaml/validation/validation_models.py b/pyaml/validation/validation_models.py index 0683d86d..1cd9e8bc 100644 --- a/pyaml/validation/validation_models.py +++ b/pyaml/validation/validation_models.py @@ -120,7 +120,7 @@ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - if getattr(cls, "validation_model", None) is not None: + if "validation_model" in cls.__dict__: raise TypeError(f"{cls.__name__} may not define validation_model manually.") cls.validation_model = cls._build_validation_model() diff --git a/tests/arrays/test_arrays.py b/tests/arrays/test_arrays.py index 001d6b14..a523f091 100644 --- a/tests/arrays/test_arrays.py +++ b/tests/arrays/test_arrays.py @@ -5,12 +5,9 @@ from pyaml.accelerator import Accelerator from pyaml.arrays.bpm import BPM -from pyaml.arrays.bpm import ConfigModel as BPMArrayConfigModel from pyaml.arrays.bpm_array import BPMArray from pyaml.arrays.cfm_magnet import CombinedFunctionMagnet -from pyaml.arrays.cfm_magnet import ConfigModel as CombinedFunctionMagnetConfigModel from pyaml.arrays.element_array import ElementArray -from pyaml.arrays.magnet import ConfigModel as MagnetArrayConfigModel from pyaml.arrays.magnet import Magnet from pyaml.arrays.magnet_array import MagnetArray @@ -228,17 +225,17 @@ def test_arrays(install_test_package): assert isinstance(magSH1AC, ElementArray) and len(magSH1AC) == 32 # Empty arrays - emptyMag = Magnet(MagnetArrayConfigModel(name="EmptyMag", elements=[])) + emptyMag = Magnet(name="EmptyMag", elements=[]) emptyMag.fill_array(sr.design) # Attach the array v = sr.design.get_magnets("EmptyMag").strengths.get() # Ensure good attach assert np.shape(v) == (0,) - emptyBPM = BPM(BPMArrayConfigModel(name="emptyBPM", elements=[])) + emptyBPM = BPM(name="emptyBPM", elements=[]) emptyBPM.fill_array(sr.design) # Attach the array v = sr.design.get_bpms("emptyBPM").positions.get() # Ensure good attach assert np.shape(v) == (0,) - emptyCFM = CombinedFunctionMagnet(CombinedFunctionMagnetConfigModel(name="emptyCFM", elements=[])) + emptyCFM = CombinedFunctionMagnet(name="emptyCFM", elements=[]) emptyCFM.fill_array(sr.design) # Attach the array v = sr.design.get_cfm_magnets("emptyCFM").strengths.get() # Ensure good attach assert np.shape(v) == (0,)