Another Pydantic v2 rewrite

.github/workflows/ci.yaml

@@ -14,14 +14,13 @@ defaults:
 
 jobs:
   test:
-    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, OpenMM ${{ matrix.openmm }}, Pydantic ${{ matrix.pydantic-version }}
+    name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }}, OpenMM ${{ matrix.openmm }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [macOS-latest, ubuntu-latest]
         python-version: ["3.10", "3.11", "3.12"]
-        pydantic-version: ["1", "2"]
         openmm: [true, false]
 
     steps:
@@ -33,7 +32,6 @@ jobs:
         environment-file: devtools/conda-envs/test_env.yaml
         create-args: >-
           python=${{ matrix.python-version }}
-          pydantic=${{ matrix.pydantic-version }}
 
     - name: Optionally install OpenMM
       if: ${{ matrix.openmm == true }}
@@ -43,11 +41,11 @@ jobs:
       run: python -m pip install -e .
 
     - name: Run mypy
-      if: ${{ matrix.python-version == 3.10 && matrix.pydantic-version == 2}}
+      if: ${{ matrix.python-version == 3.12 }}
       run: mypy -p "openff.models"
 
     - name: Run tests
-      run: pytest -v -Werror --cov=openff/models --cov-report=xml --color=yes openff/models/
+      run: pytest -v --cov=openff/models --cov-report=xml --color=yes openff/models/
 
     - name: CodeCov
       uses: codecov/codecov-action@v4

.pre-commit-config.yaml

@@ -39,4 +39,8 @@ repos:
   - id: pyupgrade
     files: ^openff
     exclude: openff/models/_version.py|setup.py
-    args: [--py39-plus]
+    args: [--py310-plus]
+- repo: https://github.com/adamchainz/blacken-docs
+  rev: 1.16.0
+  hooks:
+  - id: blacken-docs

README.md

@@ -4,69 +4,117 @@ openff-models
 [![GitHub Actions Build Status](https://github.com/openforcefield/openff-models/workflows/ci/badge.svg)](https://github.com/openforcefield/openff-models/actions?query=workflow%3Aci)
 [![codecov](https://codecov.io/gh/openforcefield/openff-models/branch/main/graph/badge.svg)](https://codecov.io/gh/openforcefield/openff-models/branch/main)
 
-
 Helper classes for Pydantic compatibility in the OpenFF stack
 
 ### Getting started
 
-```python3
-import pprint
+```python
+from pprint import pprint
 import json
 
 from openff.models.models import DefaultModel
-from openff.models.types import ArrayQuantity, FloatQuantity
-from openff.units import unit
+from openff.models.dimension_types import LengthQuantity
+from openff.models.unit_types import (
+    OnlyAMUQuantity,
+    OnlyElementaryChargeQuantity,
+)
+from openff.units import Quantity, unit
 
 
 class Atom(DefaultModel):
-    mass: FloatQuantity["atomic_mass_constant"]
-    charge: FloatQuantity["elementary_charge"]
-    some_array: ArrayQuantity["nanometer"]
+    mass: OnlyAMUQuantity
+    charge: OnlyElementaryChargeQuantity
+    some_array: LengthQuantity
 
 
 atom = Atom(
-    mass=12.011 * unit.atomic_mass_constant,
+    mass=Quantity(12.011, "atomic_mass_constant"),
     charge=0.0 * unit.elementary_charge,
-    some_array=unit.Quantity([4, -1, 0], unit.nanometer),
+    some_array=Quantity([4, -1, 0], "nanometer"),
 )
 
-print(atom.dict())
-# {'mass': <Quantity(12.011, 'atomic_mass_constant')>, 'charge': <Quantity(0.0, 'elementary_charge')>, 'some_array': <Quantity([ 4 -1  0], 'nanometer')>}
+pprint(atom.dict())
+# {'charge': <Quantity(0.0, 'elementary_charge')>,
+#  'mass': <Quantity(12.011, 'unified_atomic_mass_unit')>,
+#  'some_array': <Quantity([ 4 -1  0], 'nanometer')>}
+#
 
 # Note that unit-bearing fields use custom serialization into a dict with separate key-val pairs for
 # the unit (as a string) and unitless quantities (in whatever shape the data is)
-print(atom.json())
-# {"mass": "{\"val\": 12.011, \"unit\": \"atomic_mass_constant\"}", "charge": "{\"val\": 0.0, \"unit\": \"elementary_charge\"}", "some_array": "{\"val\": [4, -1, 0], \"unit\": \"nanometer\"}"}
+pprint(atom.json())
+# ('{"mass":"{\\"val\\": 12.011, \\"unit\\": '
+#  '\\"unified_atomic_mass_unit\\"}","charge":"{\\"val\\": 0.0, \\"unit\\": '
+#  '\\"elementary_charge\\"}","some_array":"{\\"val\\": [4, -1, 0], \\"unit\\": '
+#  '\\"nanometer\\"}"}')
 
 # The same thing, just more human-readable
-pprint.pprint(json.loads(atom.json()))
+pprint(json.loads(atom.json()))
 # {'charge': '{"val": 0.0, "unit": "elementary_charge"}',
-#  'mass': '{"val": 12.011, "unit": "atomic_mass_constant"}',
+#  'mass': '{"val": 12.011, "unit": "unified_atomic_mass_unit"}',
 #  'some_array': '{"val": [4, -1, 0], "unit": "nanometer"}'}
 
-# Can also roundtrip through these representations
+# Can also roundtrip through dict/JSON representations
 assert Atom(**atom.dict()).charge.m == 0.0
 assert Atom.parse_raw(atom.json()).charge.m == 0.0
 ```
 
-Currently, models can also be defined with a simple `unit.Quantity` annotation. This keeps serialization functionality but does not pick up the validaiton features of the custom types, i.e. dimensionality validation.
+`openff-models` ships a number of these annotated types by default, covering common dimensions and units. For those that aren't covered, there are helper classes to construct them.
 
-```python3
-import json
+```python
+from openff.units import Quantity
 
-from openff.units import unit
 from openff.models.models import DefaultModel
+from openff.models.types.dimension_types import build_dimension_type, LengthQuantity
 
+BondForceConstant = build_dimension_type("kilocalorie / mole / angstrom ** 2")
 
-class Atom(DefaultModel):
-    mass: unit.Quantity = unit.Quantity(0.0, unit.amu)
 
-json.loads(Atom(mass=12.011 * unit.atomic_mass_constant).json())
-# {'mass': '{"val": 12.011, "unit": "atomic_mass_constant"}'}
+class BondParameter(DefaultModel):
+    length: LengthQuantity
+    k: BondForceConstant
+
+
+BondParameter(
+    length=Quantity("1.35 angstrom"),
+    k=Quantity("400 kilocalorie_per_mole / angstrom **2"),
+).model_dump()
+# {'length': 1.35 <Unit('angstrom')>,
+#  'k': 400.0 <Unit('kilocalorie_per_mole / angstrom ** 2')>}
+```
+
+Currently, models can also be defined with a simple `Quantity` annotation. This keeps serialization functionality but does not pick up the validaiton features of the custom types, i.e. dimensionality or unit validation, duck-typing from `str` or other types that can be coerced into `Quantity`.
 
-# This model does have instructions to keep masses in mass units
-json.loads(Atom(mass=12.011 * unit.nanometer).json())
-# {'mass': '{"val": 12.011, "unit": "nanometer"}'}
+```python
+import json
+
+from pydantic import ValidationError
+
+from openff.units import Quantity
+from openff.models.models import DefaultModel
+
+
+class Atom(DefaultModel):
+    mass: Quantity
+
+
+# Works fine if given a Quantity
+Atom(mass=Quantity("0 amu")).model_dump()
+# {'mass': 0 <Unit('unified_atomic_mass_unit')>}
+
+# Won't automagically convert str, whereas MassQuantity or other annotated types would
+try:
+    Atom(mass="0 amu").model_dump()
+except ValidationError as error:
+    print(error)
+# ValidationError: 1 validation error for Atom
+# mass
+#   Input should be an instance of Quantity [type=is_instance_of, input_value='0 amu', input_type=str]
+#     For further information visit https://errors.pydantic.dev/2.6/v/is_instance_of
+
+# And it will gladly accept a Quantity with incompatible units,
+# which may lead to surprising results
+Atom(mass=Quantity("0 nanometer")).model_dump()
+# {'mass': 0 <Unit('nanometer')>}
 ```
 
 ### Copyright

devtools/conda-envs/test_env.yaml

@@ -4,7 +4,7 @@ channels:
 dependencies:
   - python
   - pip
-  - pydantic
+  - pydantic =2
   - openff-units
   - openff-utilities
   - pytest
@@ -14,3 +14,4 @@ dependencies:
   - unyt =3
   - pip:
     - types-setuptools
+    - git+https://github.com/openforcefield/openff-units.git@main

openff/models/_pydantic.py

@@ -1,4 +1 @@
-try:
-    from pydantic.v1 import BaseModel
-except ImportError:
-    from pydantic import BaseModel  # type: ignore[assignment]
+from pydantic import BaseModel

openff/models/_tests/integration_tests/test_models.py

@@ -0,0 +1,26 @@
+import pytest
+from openff.units import Quantity, unit
+
+from openff.models.models import DefaultModel
+from openff.models.types.dimension_types import LengthQuantity
+
+
+def test_model_pint_openmm_unit():
+    pytest.importorskip("openmm")
+    pytest.importorskip("unyt")
+
+    import openmm.unit
+    import unyt
+
+    class Bagel(DefaultModel):
+        x: LengthQuantity
+        y: LengthQuantity
+        z: LengthQuantity
+
+    bagel = Bagel(
+        x=1.0 * unit.angstrom,
+        y=1.0 * openmm.unit.angstrom,
+        z=1.0 * unyt.angstrom,
+    )
+
+    assert bagel.x == bagel.y == bagel.z == Quantity("1.0 angstrom")