0.8/pydantic2 (Sourcery refactored)

glotaran/builtin/elements/baseline/element.py

@@ -15,8 +15,8 @@
 
 class BaselineElement(Element):
     type: Literal["baseline"]  # type:ignore[assignment]
-    register_as = "baseline"  # type:ignore[pydantic-field]
-    unique = True  # type:ignore[pydantic-field]
+    register_as: str = "baseline"  # type:ignore[misc]
+    unique: bool = True
 
     def clp_label(self) -> str:
         return f"baseline_{self.label}"

glotaran/builtin/elements/clp_guide/element.py

@@ -14,8 +14,8 @@
 
 class ClpGuideElement(Element):
     type: Literal["clp-guide"]  # type:ignore[assignment]
-    register_as = "clp-guide"  # type:ignore[pydantic-field]
-    exclusive = True  # type:ignore[pydantic-field]
+    register_as: str = "clp-guide"  # type:ignore[misc]
+    exclusive: bool = True
     target: str
 
     def calculate_matrix(

glotaran/builtin/elements/coherent_artifact/element.py

@@ -13,16 +13,17 @@
 from glotaran.model import Element
 from glotaran.model import GlotaranModelError
 from glotaran.model import ParameterType
+from glotaran.model.data_model import DataModel
 
 if TYPE_CHECKING:
     from glotaran.typing.types import ArrayLike
 
 
 class CoherentArtifactElement(Element):
     type: Literal["coherent-artifact"]  # type:ignore[assignment]
-    register_as = "coherent-artifact"  # type:ignore[pydantic-field]
-    dimension = "time"  # type:ignore[pydantic-field]
-    data_model_type = ActivationDataModel  # type:ignore[pydantic-field]
+    register_as: str = "coherent-artifact"  # type:ignore[misc]
+    dimension: str = "time"
+    data_model_type: type[DataModel] = ActivationDataModel  # type:ignore[misc,valid-type]
     order: int
     width: ParameterType | None = None
 

glotaran/builtin/elements/damped_oscillation/element.py

@@ -20,6 +20,7 @@
 from glotaran.model import Element
 from glotaran.model import Item
 from glotaran.model import ParameterType
+from glotaran.model.data_model import DataModel
 
 if TYPE_CHECKING:
     from glotaran.typing.types import ArrayLike
@@ -32,9 +33,11 @@ class Oscillation(Item):
 
 class DampedOscillationElement(Element):
     type: Literal["damped-oscillation"]  # type:ignore[assignment]
-    register_as = "damped-oscillation"  # type:ignore[pydantic-field]
+    register_as: str = "damped-oscillation"  # type:ignore[misc]
     dimension: str = "time"
-    data_model_type = ActivationDataModel  # type:ignore[pydantic-field]
+    data_model_type: type[  # type:ignore[misc,valid-type]
+        DataModel
+    ] = ActivationDataModel
     oscillations: dict[str, Oscillation]
 
     def calculate_matrix(  # type:ignore[override]

glotaran/builtin/elements/kinetic/element.py

@@ -14,20 +14,21 @@
 from glotaran.builtin.items.activation import MultiGaussianActivation
 from glotaran.builtin.items.activation import add_activation_to_result_data
 from glotaran.model import ExtendableElement
+from glotaran.model.data_model import DataModel
 from glotaran.model.data_model import is_data_model_global
 
 if TYPE_CHECKING:
     from glotaran.typing.types import ArrayLike
 
 
 class KineticElement(ExtendableElement, Kinetic):
-    type: Literal["kinetic"] = Literal["kinetic"]  # type:ignore[assignment]
-    register_as = "kinetic"  # type:ignore[pydantic-field]
-    data_model_type = ActivationDataModel  # type:ignore[pydantic-field]
+    type: Literal["kinetic"]  # type:ignore[assignment]
+    register_as: str = "kinetic"  # type:ignore[misc]
+    data_model_type: type[DataModel] = ActivationDataModel  # type:ignore[misc, valid-type]
     dimension: str = "time"
 
     def extend(self, other: KineticElement):  # type:ignore[override]
-        return other.copy(update={"rates": self.rates | other.rates})
+        return other.model_copy(update={"rates": self.rates | other.rates})
 
     # TODO: consolidate parent method.
     @classmethod
@@ -48,7 +49,11 @@ def combine(cls, kinetics: list[KineticElement]) -> KineticElement:  # type:igno
             The combined KMatrix.
 
         """
-        return cls(rates=reduce(lambda lhs, rhs: lhs | rhs, [k.rates for k in kinetics]), label="")
+        return cls(
+            type="kinetic",
+            rates=reduce(lambda lhs, rhs: lhs | rhs, [k.rates for k in kinetics]),
+            label="",
+        )
 
     @staticmethod
     def combine_matrices(lhs: ArrayLike, rhs: ArrayLike) -> ArrayLike:

glotaran/builtin/elements/spectral/element.py

@@ -21,9 +21,9 @@ class SpectralDataModel(DataModel):
 
 class SpectralElement(Element):
     type: Literal["spectral"]  # type:ignore[assignment]
+    register_as: str = "spectral"  # type:ignore[misc]
     dimension: str = "spectral"
-    register_as = "spectral"  # type:ignore[pydantic-field]
-    data_model_type = SpectralDataModel  # type:ignore[pydantic-field]
+    data_model_type: type[DataModel] = SpectralDataModel  # type:ignore[misc,valid-type]
     shapes: dict[str, SpectralShape.get_annotated_type()]  # type:ignore[valid-type]
 
     def calculate_matrix(  # type:ignore[override]

glotaran/io/preprocessor/preprocessor.py

@@ -6,15 +6,13 @@
 
 import xarray as xr
 from pydantic import BaseModel
+from pydantic import ConfigDict
 
 
 class PreProcessor(BaseModel, abc.ABC):
     """A base class for pre=processors."""
 
-    class Config:
-        """Config for BaseModel."""
-
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @abc.abstractmethod
     def apply(self, data: xr.DataArray) -> xr.DataArray:

glotaran/io/preprocessor/test/test_preprocessor.py

@@ -35,11 +35,15 @@ def test_to_from_dict():
         .correct_baseline_value(1)
         .correct_baseline_average({"dim_1": slice(0, 2)})
     )
-    pl_dict = pl.dict()
+    pl_dict = pl.model_dump()
     assert pl_dict == {
         "actions": [
             {"action": "baseline-value", "value": 1.0},
-            {"action": "baseline-average", "select": {"dim_1": slice(0, 2)}, "exclude": None},
+            {
+                "action": "baseline-average",
+                "select": {"dim_1": slice(0, 2)},
+                "exclude": None,
+            },
         ]
     }
-    assert PreProcessingPipeline.parse_obj(pl_dict) == pl
+    assert PreProcessingPipeline.model_validate(pl_dict) == pl

glotaran/model/data_model.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Generator
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Literal
 from uuid import uuid4
@@ -22,6 +23,9 @@
 from glotaran.model.weight import Weight
 from glotaran.parameter import Parameters
 
+if TYPE_CHECKING:
+    from glotaran.project.library import ModelLibrary
+
 
 class ExclusiveModelIssue(ItemIssue):
     """Issue for exclusive elements."""
@@ -202,12 +206,17 @@ class DataModel(Item):
     def create_class_for_elements(elements: set[type[Element]]) -> type[DataModel]:
         data_model_cls_name = f"GlotaranDataModel_{str(uuid4()).replace('-','_')}"
         data_models = tuple(
-            {e.data_model_type for e in elements if e.data_model_type is not None}
+            {
+                e.model_fields["data_model_type"].default
+                for e in elements
+                if "data_model_type" in e.model_fields
+                and e.model_fields["data_model_type"].default is not None
+            }
         ) + (DataModel,)
         return create_model(data_model_cls_name, __base__=data_models)
 
     @classmethod
-    def from_dict(cls, library: dict[str, Element], model_dict: dict[str, Any]) -> DataModel:
+    def from_dict(cls, library: ModelLibrary, model_dict: dict[str, Any]) -> DataModel:
         element_labels = model_dict.get("elements", []) + model_dict.get("global_elements", [])
         if len(element_labels) == 0:
             raise GlotaranModelError("No element defined for dataset")
@@ -314,11 +323,11 @@ def iterate_data_model_global_elements(
 
 def resolve_data_model(
     model: DataModel,
-    library: dict[str, Element],
+    library: ModelLibrary,
     parameters: Parameters,
     initial: Parameters | None = None,
 ) -> DataModel:
-    model = model.copy()
+    model = model.model_copy()
     model.elements = [library[m] if isinstance(m, str) else m for m in model.elements]
     if model.global_elements is not None:
         model.global_elements = [

glotaran/model/experiment_model.py

@@ -1,35 +1,34 @@
 """This module contains the dataset group."""
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
 from typing import Any
 from typing import Literal
 
 from pydantic import BaseModel
-from pydantic import Extra
+from pydantic import ConfigDict
 from pydantic import Field
 
 from glotaran.model.clp_constraint import ClpConstraint
 from glotaran.model.clp_penalties import EqualAreaPenalty
 from glotaran.model.clp_relation import ClpRelation
 from glotaran.model.data_model import DataModel
 from glotaran.model.data_model import resolve_data_model
-from glotaran.model.element import Element
 from glotaran.model.errors import ItemIssue
 from glotaran.model.item import ParameterType
 from glotaran.model.item import get_item_issues
 from glotaran.model.item import resolve_item_parameters
 from glotaran.model.item import resolve_parameter
 from glotaran.parameter import Parameters
 
+if TYPE_CHECKING:
+    from glotaran.project.library import ModelLibrary
+
 
 class ExperimentModel(BaseModel):
     """A dataset group for optimization."""
 
-    class Config:
-        """Config for pydantic.BaseModel."""
-
-        arbitrary_types_allowed = True
-        extra = Extra.forbid
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
 
     clp_link_tolerance: float = 0.0
     clp_link_method: Literal["nearest", "backward", "forward"] = "nearest"
@@ -45,24 +44,25 @@ class Config:
         "variable_projection", description="The residual function to use."
     )
     scale: dict[str, ParameterType] = Field(
-        default_factory=dict, description="The scales of of the datasets in the experiment."
+        default_factory=dict,
+        description="The scales of of the datasets in the experiment.",
     )
 
     @classmethod
-    def from_dict(cls, library: dict[str, Element], model_dict: dict[str, Any]) -> ExperimentModel:
+    def from_dict(cls, library: ModelLibrary, model_dict: dict[str, Any]) -> ExperimentModel:
         model_dict["datasets"] = {
             label: DataModel.from_dict(library, dataset)
             for label, dataset in model_dict.get("datasets", {}).items()
         }
-        return cls.parse_obj(model_dict)
+        return cls.model_validate(model_dict)
 
     def resolve(
         self,
-        library: dict[str, Element],
+        library: ModelLibrary,
         parameters: Parameters,
         initial: Parameters | None = None,
     ) -> ExperimentModel:
-        result = self.copy()
+        result = self.model_copy()
         result.datasets = {
             label: resolve_data_model(dataset, library, parameters, initial)
             for label, dataset in self.datasets.items()