Update GMSO to work with pydantic 2.0

gmso/abc/abstract_connection.py

@@ -1,6 +1,6 @@
 from typing import Optional, Sequence
 
-from pydantic import Field, root_validator
+from pydantic import ConfigDict, Field, model_validator
 
 from gmso.abc.abstract_site import Site
 from gmso.abc.gmso_base import GMSOBase
@@ -65,7 +65,8 @@ def _get_members_types_or_classes(self, to_return):
             ]
             return tc if all(tc) else None
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def validate_fields(cls, values):
         connection_members = values.get("connection_members")
 
@@ -101,10 +102,12 @@ def __repr__(self):
     def __str__(self):
         return f"<{self.__class__.__name__} {self.name}, id: {id(self)}> "
 
-    class Config:
-        fields = {"name_": "name", "connection_members_": "connection_members"}
-
-        alias_to_fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        fields={"name_": "name", "connection_members_": "connection_members"},
+        alias_to_fields={
             "name": "name_",
             "connection_members": "connection_members_",
-        }
+        },
+    )

gmso/abc/abstract_potential.py

@@ -2,7 +2,7 @@
 from abc import abstractmethod
 from typing import Any, Dict, Iterator, List
 
-from pydantic import Field, validator
+from pydantic import ConfigDict, Field, field_validator
 
 from gmso.abc.gmso_base import GMSOBase
 from gmso.utils.expression import PotentialExpression
@@ -115,7 +115,8 @@ def delete_tag(self, tag: str) -> None:
     def pop_tag(self, tag: str) -> Any:
         return self.tags.pop(tag, None)
 
-    @validator("potential_expression_", pre=True)
+    @field_validator("potential_expression_", mode="before")
+    @classmethod
     def validate_potential_expression(cls, v):
         if isinstance(v, dict):
             v = PotentialExpression(**v)
@@ -154,17 +155,17 @@ def __str__(self):
             f"id: {id(self)}>"
         )
 
-    class Config:
-        """Pydantic configuration for the potential objects."""
-
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        fields={
             "name_": "name",
             "potential_expression_": "potential_expression",
             "tags_": "tags",
-        }
-
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "name": "name_",
             "potential_expression": "potential_expression_",
             "tags": "tags_",
-        }
+        },
+    )

gmso/abc/abstract_site.py

@@ -4,7 +4,14 @@
 
 import numpy as np
 import unyt as u
-from pydantic import Field, StrictInt, StrictStr, validator
+from pydantic import (
+    ConfigDict,
+    Field,
+    StrictInt,
+    StrictStr,
+    field_validator,
+    validator,
+)
 from unyt.exceptions import InvalidUnitOperation
 
 from gmso.abc.gmso_base import GMSOBase
@@ -120,7 +127,8 @@ def __str__(self):
             f"label: {self.label if self.label else None} id: {id(self)}>"
         )
 
-    @validator("position_")
+    @field_validator("position_")
+    @classmethod
     def is_valid_position(cls, position):
         """Validate attribute position."""
         if position is None:
@@ -148,6 +156,8 @@ def is_valid_position(cls, position):
 
         return position
 
+    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
     @validator("name_", pre=True, always=True)
     def inject_name(cls, value):
         if value == "" or value is None:
@@ -162,27 +172,25 @@ def __new__(cls, *args: Any, **kwargs: Any) -> SiteT:
         else:
             return object.__new__(cls)
 
-    class Config:
-        """Pydantic configuration for site objects."""
-
-        arbitrary_types_allowed = True
-
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        fields={
             "name_": "name",
             "position_": "position",
             "label_": "label",
             "group_": "group",
             "molecule_": "molecule",
             "residue_": "residue",
-        }
-
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "name": "name_",
             "position": "position_",
             "label": "label_",
             "group": "group_",
             "molecule": "molecule_",
             "residue": "residue_",
-        }
-
-        validate_assignment = True
+        },
+        validate_assignment=True,
+    )

gmso/abc/gmso_base.py

@@ -4,13 +4,14 @@
 from abc import ABC
 from typing import Any, ClassVar, Type
 
-from pydantic import BaseModel
-from pydantic.validators import dict_validator
+from pydantic import BaseModel, ConfigDict, validators
 
 from gmso.abc import GMSOJSONHandler
 from gmso.abc.auto_doc import apply_docs
 from gmso.abc.serialization_utils import dict_to_unyt
 
+dict_validator = validators.getattr_migration("dict_validator")
+
 
 class GMSOBase(BaseModel, ABC):
     """A BaseClass to all abstract classes in GMSO."""
@@ -116,11 +117,12 @@ def __get_validators__(cls) -> "CallableGenerator":
         """Get the validators of the object."""
         yield cls.validate
 
-    class Config:
-        """Pydantic configuration for base object."""
-
-        arbitrary_types_allowed = True
-        alias_to_fields = dict()
-        extra = "forbid"
-        json_encoders = GMSOJSONHandler.json_encoders
-        allow_population_by_field_name = True
+    # TODO[pydantic]: The following keys were removed: `json_encoders`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        alias_to_fields=dict(),
+        extra="forbid",
+        json_encoders=GMSOJSONHandler.json_encoders,
+        populate_by_name=True,
+    )

gmso/core/angle.py

@@ -1,7 +1,7 @@
 """Support for 3-partner connections between gmso.core.Atoms."""
 from typing import Callable, ClassVar, Optional, Tuple
 
-from pydantic import Field
+from pydantic import ConfigDict, Field
 
 from gmso.abc.abstract_connection import Connection
 from gmso.core.angle_type import AngleType
@@ -83,16 +83,17 @@ def __setattr__(self, key, value):
         else:
             super(Angle, self).__setattr__(key, value)
 
-    class Config:
-        """Support pydantic configuration for attributes and behavior."""
-
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        fields={
             "connection_members_": "connection_members",
             "angle_type_": "angle_type",
             "restraint_": "restraint",
-        }
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "connection_members": "connection_members_",
             "angle_type": "angle_type_",
             "restraint": "restraint_",
-        }
+        },
+    )

gmso/core/angle_type.py

@@ -1,7 +1,7 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
+from pydantic import ConfigDict, Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
@@ -77,13 +77,15 @@ def member_types(self):
     def member_classes(self):
         return self.__dict__.get("member_classes_")
 
-    class Config:
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        fields={
             "member_types_": "member_types",
             "member_classes_": "member_classes",
-        }
-
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "member_types": "member_types_",
             "member_classes": "member_classes_",
-        }
+        },
+    )

gmso/core/atom.py

@@ -3,7 +3,7 @@
 from typing import Optional, Union
 
 import unyt as u
-from pydantic import Field, validator
+from pydantic import ConfigDict, Field, field_validator
 
 from gmso.abc.abstract_site import Site
 from gmso.core.atom_type import AtomType
@@ -122,7 +122,8 @@ def __lt__(self, other):
                 f"Cannot compare equality between {type(self)} and {type(other)}"
             )
 
-    @validator("charge_")
+    @field_validator("charge_")
+    @classmethod
     def is_valid_charge(cls, charge):
         """Ensure that the charge is physically meaningful."""
         if charge is None:
@@ -137,7 +138,8 @@ def is_valid_charge(cls, charge):
 
         return charge
 
-    @validator("mass_")
+    @field_validator("mass_")
+    @classmethod
     def is_valid_mass(cls, mass):
         """Ensure that the mass is physically meaningful."""
         if mass is None:
@@ -150,23 +152,21 @@ def is_valid_mass(cls, mass):
             ensure_valid_dimensions(mass, default_mass_units)
         return mass
 
-    class Config:
-        """Pydantic configuration for the atom class."""
-
-        extra = "forbid"
-
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        extra="forbid",
+        fields={
             "charge_": "charge",
             "mass_": "mass",
             "element_": "element",
             "atom_type_": "atom_type",
-        }
-
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "charge": "charge_",
             "mass": "mass_",
             "element": "element_",
             "atom_type": "atom_type_",
-        }
-
-        validate_assignment = True
+        },
+        validate_assignment=True,
+    )

gmso/core/atom_type.py

@@ -3,7 +3,7 @@
 from typing import Optional, Set
 
 import unyt as u
-from pydantic import Field, validator
+from pydantic import ConfigDict, Field, field_validator
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils._constants import UNIT_WARNING_STRING
@@ -186,7 +186,8 @@ def __repr__(self):
         )
         return desc
 
-    @validator("mass_", pre=True)
+    @field_validator("mass_", mode="before")
+    @classmethod
     def validate_mass(cls, mass):
         """Check to see that a mass is a unyt array of the right dimension."""
         default_mass_units = u.gram / u.mol
@@ -198,7 +199,8 @@ def validate_mass(cls, mass):
 
         return mass
 
-    @validator("charge_", pre=True)
+    @field_validator("charge_", mode="before")
+    @classmethod
     def validate_charge(cls, charge):
         """Check to see that a charge is a unyt array of the right dimension."""
         if not isinstance(charge, u.unyt_array):
@@ -222,25 +224,25 @@ def _default_potential_expr():
             },
         )
 
-    class Config:
-        """Pydantic configuration of the attributes for an atom_type."""
-
-        fields = {
+    # TODO[pydantic]: The following keys were removed: `fields`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(
+        fields={
             "mass_": "mass",
             "charge_": "charge",
             "atomclass_": "atomclass",
             "overrides_": "overrides",
             "doi_": "doi",
             "description_": "description",
             "definition_": "definition",
-        }
-
-        alias_to_fields = {
+        },
+        alias_to_fields={
             "mass": "mass_",
             "charge": "charge_",
             "atomclass": "atomclass_",
             "overrides": "overrides_",
             "doi": "doi_",
             "description": "description_",
             "definition": "definition_",
-        }
+        },
+    )