Make config models into schemas

pyaml/accelerator.py

@@ -4,23 +4,20 @@
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from .arrays.array import ArrayConfig
-from .common.element import Element
+from .arrays.array import Array, ArraySchema
+from .common.element import Element, ElementSchema
 from .common.element_holder import ElementHolder
 from .common.exception import PyAMLConfigException
 from .configuration import ConfigurationManager, UnsupportedConfigurationRootError
 from .configuration.factory import Factory
-from .control.controlsystem import ControlSystem
-from .lattice.simulator import Simulator
+from .control.controlsystem import ControlSystem, ControlSystemSchema
+from .lattice.simulator import Simulator, SimulatorSchema
 from .yellow_pages import YellowPages
 
-# Define the main class name for this module
-PYAMLCLASS = "Accelerator"
 
-
-class ConfigModel(BaseModel):
+class AcceleratorSchema(BaseModel):
     """
-    Configuration model for Accelerator
+    Schema for Accelerator
 
     Parameters
     ----------
@@ -50,83 +47,99 @@ class ConfigModel(BaseModel):
         Element list
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+    model_config = ConfigDict(extra="forbid")
 
     facility: str
     machine: str
     energy: float
     alphac: float | None = None
     harmonic_number: int | None = None
-    controls: list[ControlSystem] = None
-    simulators: list[Simulator] = None
+    controls: list[ControlSystemSchema] = None
+    simulators: list[SimulatorSchema] = None
     data_folder: str
     description: str | None = None
-    arrays: list[ArrayConfig] = Field(default=None, repr=False)
-    devices: list[Element] = Field(repr=False)
+    arrays: list[ArraySchema] = Field(default=None, repr=False)
+    devices: list[ElementSchema] = Field(default=None, repr=False)
 
 
 class Accelerator(object):
     """PyAML top level class"""
 
-    def __init__(self, cfg: ConfigModel):
-        self._cfg = cfg
+    def __init__(
+        self,
+        facility: str,
+        machine: str,
+        energy: float,
+        data_folder: str,
+        alphac: float | None = None,
+        harmonic_number: int | None = None,
+        controls: list[ControlSystem] = None,
+        simulators: list[Simulator] = None,
+        description: str | None = None,
+        arrays: list[Array] | None = None,
+        devices: list[Element] | None = None,
+    ):
+        self.facility = facility
+        self.machine = machine
+
         __design = None
         __live = None
+
         self._controls: dict[str, ElementHolder] = {}
         self._simulators: dict[str, ElementHolder] = {}
 
-        if cfg.controls is not None:
-            for c in cfg.controls:
+        if controls is not None:
+            for c in controls:
                 if c.name() == "live":
                     self.__live = c
                 else:
                     # Add as dynamic attribute
                     setattr(self, c.name(), c)
-                c.fill_device(cfg.devices)
+                c.fill_device(devices)
                 c._peer = self
                 self._controls[c.name()] = c
 
-        if cfg.simulators is not None:
-            for s in cfg.simulators:
+        if simulators is not None:
+            for s in simulators:
                 if s.name() == "design":
                     self.__design = s
                 else:
                     # Add as dynamic attribute
                     setattr(self, s.name(), s)
-                s.fill_device(cfg.devices)
+                s.fill_device(devices)
                 s._peer = self
                 self._simulators[s.name()] = s
 
-        if cfg.arrays is not None:
-            for a in cfg.arrays:
-                if cfg.simulators is not None:
-                    for s in cfg.simulators:
+        if arrays is not None:
+            for a in arrays:
+                if simulators is not None:
+                    for s in simulators:
                         a.fill_array(s)
-                if cfg.controls is not None:
-                    for c in cfg.controls:
+                if controls is not None:
+                    for c in controls:
                         a.fill_array(c)
 
-        if cfg.energy is not None:
-            self.set_energy(cfg.energy)
+        if energy is not None:
+            self.set_energy(energy)
 
-        if cfg.alphac is not None:
-            self.set_mcf(cfg.alphac)
+        if alphac is not None:
+            self.set_mcf(alphac)
 
-        if cfg.harmonic_number is not None:
-            self.set_harmonic_number(cfg.harmonic_number)
+        if harmonic_number is not None:
+            self.set_harmonic_number(harmonic_number)
 
         self._yellow_pages = YellowPages(self)
 
         self.post_init()
 
     def _set_properties(self, method: str, value):
         # Sets global property
-        if self._cfg.simulators is not None:
-            for s in self._cfg.simulators:
+        if self._simulators is not None:
+            for s in self._simulators:
                 m = getattr(s, method)
                 m(value)
-        if self._cfg.controls is not None:
-            for c in self._cfg.controls:
+        if self._controls is not None:
+            for c in self._controls:
                 m = getattr(c, method)
                 m(value)
 
@@ -180,24 +193,24 @@ def add_device(self, config: dict, ignore_external=False):
                 "Invalid device type, Element or sub classes of Element expected " + f"but got {dev.__class__.__name__}"
             )
 
-        self._cfg.devices.append(dev)
-        if self._cfg.controls is not None:
-            for c in self._cfg.controls:
+        self._devices.append(dev)
+        if self._controls is not None:
+            for c in self._controls:
                 c.fill_device([dev])
 
-        if self._cfg.simulators is not None:
-            for s in self._cfg.simulators:
+        if self._simulators is not None:
+            for s in self._simulators:
                 s.fill_device([dev])
 
     def post_init(self):
         """
         Method triggered after all initialisations are done
         """
-        if self._cfg.simulators is not None:
-            for s in self._cfg.simulators:
+        if self._simulators is not None:
+            for s in self._simulators:
                 s.post_init()
-        if self._cfg.controls is not None:
-            for c in self._cfg.controls:
+        if self._controls is not None:
+            for c in self._controls:
                 c.post_init()
 
     def get_description(self) -> str:

pyaml/arrays/array.py

@@ -1,5 +1,5 @@
 """
-Array configuration
+Base classes for arrays
 """
 
 from pydantic import BaseModel, ConfigDict
@@ -9,33 +9,36 @@
 from ..common.element_holder import ElementHolder
 
 
-class ArrayConfigModel(BaseModel):
+class ArraySchema(BaseModel):
     """
-    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`.
+    Schema for configuration of array.
 
     Parameters
     ----------
     name : str
         Name of the array
     elements : list[str]
-        List of pyaml element names
+        List of element names
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+    model_config = ConfigDict(extra="forbid")
 
     name: str
     elements: list[str]
 
 
-class ArrayConfig(object):
+class Array:
     """
     Base class that implements configuration for access to arrays (families)
     """
 
-    def __init__(self, cfg: ArrayConfigModel):
-        self._cfg = cfg
+    def __init__(
+        self,
+        name: str,
+        elements: list[str],
+    ):
+        self._name = name
+        self._elements = elements
 
     def fill_array(self, holder: ElementHolder):
         """
@@ -56,7 +59,7 @@ 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__)
+    # def __repr__(self):
+    #     # ArrayConfigModel is a super class
+    #     # ConfigModel is expected from sub classes
+    #     return repr(self._cfg).replace("ConfigModel", self.__class__.__name__)

pyaml/arrays/bpm.py

@@ -1,15 +1,8 @@
 from ..common.element_holder import ElementHolder
-from .array import ArrayConfig, ArrayConfigModel
+from .array import Array
 
-# Define the main class name for this module
-PYAMLCLASS = "BPM"
 
-
-class ConfigModel(ArrayConfigModel):
-    """Configuration model for :py:class:`.BPMArray`."""
-
-
-class BPM(ArrayConfig):
+class BPM(Array):
     """
     :py:class:`.BPMArray` configuration.
 
@@ -27,21 +20,15 @@ class BPM(ArrayConfig):
             - BPM_C04-02
             - BPM_C04-03
 
-    A :py:class:`.BPMArray` configuration can also be created by code using the
-    following example:
-
-    .. code-block:: python
-
-        >>>  from pyaml.arrays.bpm import BPM,ConfigModel as BPMArrayConfigModel
-        >>>  bpm_cfg = BPM(BPMArrayConfigModel(
-                        name="BPM",
-                        elements=["BPM_C04-01","BPM_C04-02","BPM_C04-03"]
-                       ))
-
     """
 
-    def __init__(self, cfg: ArrayConfigModel):
-        super().__init__(cfg)
+    def __init__(
+        self,
+        name: str,
+        elements: list[str],
+    ):
+        self._name = name
+        self._elements = elements
 
     def fill_array(self, holder: ElementHolder):
         """
@@ -66,4 +53,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)

pyaml/arrays/cfm_magnet.py

@@ -1,35 +1,20 @@
 from ..common.element_holder import ElementHolder
-from .array import ArrayConfig, ArrayConfigModel
+from .array import Array
 
-# Define the main class name for this module
-PYAMLCLASS = "CombinedFunctionMagnet"
 
-
-class ConfigModel(ArrayConfigModel):
-    """Configuration model for Combined Function Magnet array."""
-
-    ...
-
-
-class CombinedFunctionMagnet(ArrayConfig):
+class CombinedFunctionMagnet(Array):
     """
-    Combined function magnet array confirguration
-
-    Example
-    -------
-
-    A magnet array configuration can also be created by code using
-    the following example::
+    Combined function magnet array confirguration.
 
-        from pyaml.arrays.cfm_magnet import CombinedFunctionMagnet
-        from pyaml.arrays.cfm_magnet import ConfigModel as CFMagnetConfigModel
-        magArray = CombinedFunctionMagnet(
-                     CFMagnetConfigModel(name="myCFM", elements=["mag1","mag2"])
-                   )
     """
 
-    def __init__(self, cfg: ArrayConfigModel):
-        super().__init__(cfg)
+    def __init__(
+        self,
+        name: str,
+        elements: list[str],
+    ):
+        self._name = name
+        self._elements = elements
 
     def fill_array(self, holder: ElementHolder):
         """
@@ -40,4 +25,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)