Support Pydantic v1 and v2 (mosdef-hub#752)

Co-authored-by: Co Quach <[email protected]>

.github/workflows/CI.yaml

@@ -20,6 +20,7 @@ jobs:
       matrix:
         os: [macOS-latest, ubuntu-latest]
         python-version: ["3.8", "3.9", "3.10", "3.11"]
+        pydantic-version: ["1", "2"]
 
     defaults:
       run:
@@ -33,7 +34,9 @@ jobs:
         uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: environment-dev.yml
-          python-version: python=${{ matrix.python-version }}
+          create-args: >-
+            python=${{ matrix.python-version }}
+            pydantic=${{ matrix.pydantic-version }}
 
       - name: Install Package
         run: python -m pip install -e .

environment-dev.yml

@@ -9,7 +9,7 @@ dependencies:
   - unyt<=2.9.2
   - boltons
   - lxml
-  - pydantic=1.10.11
+  - pydantic
   - networkx
   - pytest
   - mbuild>=0.11.0

environment.yml

@@ -9,7 +9,7 @@ dependencies:
   - unyt<=2.9.2
   - boltons
   - lxml
-  - pydantic=1.10.11
+  - pydantic
   - networkx
   - ele>=0.2.0
   - foyer>=0.11.3

gmso/abc/abstract_connection.py

@@ -1,11 +1,14 @@
 from typing import Optional, Sequence
 
-from pydantic import Field, root_validator
-
 from gmso.abc.abstract_site import Site
 from gmso.abc.gmso_base import GMSOBase
 from gmso.exceptions import GMSOError
 
+try:
+    from pydantic.v1 import Field, root_validator
+except ImportError:
+    from pydantic import Field, root_validator
+
 
 class Connection(GMSOBase):
     __base_doc__ = """An abstract class that stores data about connections between sites.

gmso/abc/abstract_potential.py

@@ -2,11 +2,14 @@
 from abc import abstractmethod
 from typing import Any, Dict, Iterator, List
 
-from pydantic import Field, validator
-
 from gmso.abc.gmso_base import GMSOBase
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field, validator
+except ImportError:
+    from pydantic import Field, validator
+
 
 class AbstractPotential(GMSOBase):
     __base_doc__ = """An abstract potential class.

gmso/abc/abstract_site.py

@@ -4,12 +4,16 @@
 
 import numpy as np
 import unyt as u
-from pydantic import Field, StrictInt, StrictStr, validator
 from unyt.exceptions import InvalidUnitOperation
 
 from gmso.abc.gmso_base import GMSOBase
 from gmso.exceptions import GMSOError
 
+try:
+    from pydantic.v1 import Field, StrictInt, StrictStr, validator
+except ImportError:
+    from pydantic import Field, StrictInt, StrictStr, validator
+
 PositionType = Union[Sequence[float], np.ndarray, u.unyt_array]
 MoleculeType = NamedTuple("Molecule", name=StrictStr, number=StrictInt)
 ResidueType = NamedTuple("Residue", name=StrictStr, number=StrictInt)

gmso/abc/auto_doc.py

@@ -4,7 +4,10 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-from pydantic import BaseModel
+try:
+    from pydantic.v1 import BaseModel
+except ImportError:
+    from pydantic import BaseModel
 
 BASE_DOC_ATTR = "__base_doc__"
 FIELDS_IN_DOCSTRING = "__alias_to_fields__"

gmso/abc/gmso_base.py

@@ -4,13 +4,17 @@
 from abc import ABC
 from typing import Any, ClassVar, Type
 
-from pydantic import BaseModel
-from pydantic.validators import dict_validator
-
 from gmso.abc import GMSOJSONHandler
 from gmso.abc.auto_doc import apply_docs
 from gmso.abc.serialization_utils import dict_to_unyt
 
+try:
+    from pydantic.v1 import BaseModel
+    from pydantic.v1.validators import dict_validator
+except ImportError:
+    from pydantic import BaseModel
+    from pydantic.validators import dict_validator
+
 
 class GMSOBase(BaseModel, ABC):
     """A BaseClass to all abstract classes in GMSO."""

gmso/core/angle.py

@@ -1,12 +1,15 @@
 """Support for 3-partner connections between gmso.core.Atoms."""
 from typing import Callable, ClassVar, Optional, Tuple
 
-from pydantic import Field
-
 from gmso.abc.abstract_connection import Connection
 from gmso.core.angle_type import AngleType
 from gmso.core.atom import Atom
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class Angle(Connection):
     __base_doc__ = """A 3-partner connection between Atoms.

gmso/core/angle_type.py

@@ -1,11 +1,15 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class AngleType(ParametricPotential):
     __base_doc__ = """A descripton of the interaction between 3 bonded partners.

gmso/core/atom.py

@@ -3,14 +3,18 @@
 from typing import Optional, Union
 
 import unyt as u
-from pydantic import Field, validator
 
 from gmso.abc.abstract_site import Site
 from gmso.core.atom_type import AtomType
 from gmso.core.element import Element
 from gmso.utils._constants import UNIT_WARNING_STRING
 from gmso.utils.misc import ensure_valid_dimensions
 
+try:
+    from pydantic.v1 import Field, validator
+except ImportError:
+    from pydantic import Field, validator
+
 
 class Atom(Site):
     __base_doc__ = """An atom represents a single element association in a topology.

gmso/core/atom_type.py

@@ -3,7 +3,6 @@
 from typing import Optional, Set
 
 import unyt as u
-from pydantic import Field, validator
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils._constants import UNIT_WARNING_STRING
@@ -14,6 +13,11 @@
     unyt_to_hashable,
 )
 
+try:
+    from pydantic.v1 import Field, validator
+except ImportError:
+    from pydantic import Field, validator
+
 
 class AtomType(ParametricPotential):
     __base_doc__ = """A description of non-bonded interactions between sites.

gmso/core/bond.py

@@ -1,12 +1,15 @@
 """Module for 2-partner connections between sites."""
 from typing import Callable, ClassVar, Optional, Tuple
 
-from pydantic import Field
-
 from gmso.abc.abstract_connection import Connection
 from gmso.core.atom import Atom
 from gmso.core.bond_type import BondType
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class Bond(Connection):
     __base_doc__ = """A 2-partner connection between sites.

gmso/core/bond_type.py

@@ -2,11 +2,15 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class BondType(ParametricPotential):
     __base_doc__ = """A descripton of the interaction between 2 bonded partners.

gmso/core/dihedral.py

@@ -1,11 +1,14 @@
 from typing import Callable, ClassVar, Optional, Tuple
 
-from pydantic import Field
-
 from gmso.abc.abstract_connection import Connection
 from gmso.core.atom import Atom
 from gmso.core.dihedral_type import DihedralType
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class Dihedral(Connection):
     __base_doc__ = """A 4-partner connection between sites.

gmso/core/dihedral_type.py

@@ -1,11 +1,15 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class DihedralType(ParametricPotential):
     __base_doc__ = """A descripton of the interaction between 4 bonded partners.

gmso/core/element.py

@@ -6,12 +6,16 @@
 import numpy as np
 import unyt as u
 from pkg_resources import resource_filename
-from pydantic import Field
 
 from gmso.abc.gmso_base import GMSOBase
 from gmso.exceptions import GMSOError
 from gmso.utils.misc import unyt_to_hashable
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 exported = [
     "element_by_mass",
     "element_by_symbol",

gmso/core/improper.py

@@ -1,12 +1,15 @@
 """Support for improper style connections (4-member connection)."""
 from typing import Callable, ClassVar, Optional, Tuple
 
-from pydantic import Field
-
 from gmso.abc.abstract_connection import Connection
 from gmso.core.atom import Atom
 from gmso.core.improper_type import ImproperType
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class Improper(Connection):
     __base_doc__ = """sA 4-partner connection between sites.

gmso/core/improper_type.py

@@ -2,11 +2,15 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class ImproperType(ParametricPotential):
     __base_doc__ = """A description of the interaction between 4 bonded partners.

gmso/core/pairpotential_type.py

@@ -1,11 +1,15 @@
 from typing import Optional, Tuple
 
 import unyt as u
-from pydantic import Field
 
 from gmso.core.parametric_potential import ParametricPotential
 from gmso.utils.expression import PotentialExpression
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class PairPotentialType(ParametricPotential):
     __base_doc__ = """A description of custom pairwise potential between 2 AtomTypes that does not follow combination rule.

gmso/lib/potential_templates.py

@@ -5,7 +5,6 @@
 
 import sympy
 import unyt as u
-from pydantic import Field, validator
 
 from gmso.abc.abstract_potential import AbstractPotential
 from gmso.exceptions import (
@@ -16,6 +15,11 @@
 from gmso.utils.expression import PotentialExpression
 from gmso.utils.singleton import Singleton
 
+try:
+    from pydantic.v1 import Field, validator
+except ImportError:
+    from pydantic import Field, validator
+
 POTENTIAL_JSONS = list(Path(__file__).parent.glob("jsons/*.json"))
 JSON_DIR = Path.joinpath(Path(__file__).parent, "jsons")
 

gmso/parameterization/topology_parameterizer.py

@@ -5,7 +5,6 @@
 
 import networkx as nx
 from boltons.setutils import IndexedSet
-from pydantic import Field
 
 from gmso.abc.gmso_base import GMSOBase
 from gmso.core.forcefield import ForceField
@@ -29,6 +28,11 @@
 )
 from gmso.parameterization.utils import POTENTIAL_GROUPS
 
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
 
 class ParameterizationError(GMSOError):
     """Raise when parameterization fails."""

gmso/tests/test_angle.py

@@ -1,5 +1,4 @@
 import pytest
-from pydantic import ValidationError
 
 from gmso.core.angle import Angle
 from gmso.core.angle_type import AngleType
@@ -8,6 +7,11 @@
 from gmso.core.topology import Topology
 from gmso.tests.base_test import BaseTest
 
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:
+    from pydantic import ValidationError
+
 
 class TestAngle(BaseTest):
     def test_angle_nonparametrized(self):