Define Pydantic schema on Quantity

.github/workflows/ci.yaml

@@ -32,6 +32,9 @@ jobs:
           - py312pint25openmm
           - py313pint25openmm
           - py314pint25openmm
+          - py312pydantic
+          - py313pydantic
+          - py314pydantic
 
     steps:
     - uses: actions/checkout@v7

.github/workflows/ci_pypi.yaml

@@ -14,7 +14,7 @@ defaults:
 
 jobs:
   test:
-    name: 📦=${{ matrix.install-source }}, 💻=${{ matrix.os }}, 🐍=${{ matrix.python-version }}, 🍺=${{ matrix.pint-version }}, ⚛️=${{ matrix.openmm }}
+    name: 📦=${{ matrix.install-source }}, 💻=${{ matrix.os }}, 🐍=${{ matrix.python-version }}, 🍺=${{ matrix.pint-version }}, ⚛️==${{ matrix.openmm }}, 🤓=${{ matrix.pydantic }}
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
@@ -23,6 +23,7 @@ jobs:
         openmm: ["true", "false"]
         python-version: ["3.12", "3.13", "3.14"]
         pint-version: ["0.24", "0.25"]
+        pydantic: ["true", "false"]
 
     env:
       PYTEST_ARGS: -v -n auto
@@ -72,6 +73,10 @@ jobs:
             pint~=${{ matrix.pint-version }} \
             -e ."[test,typing]" downstream_dummy/
 
+    - name: Optionally install Pydantic from PyPI
+      if: ${{ matrix.install-source == 'pypi' && matrix.pydantic == 'true' }}
+      run: pip install "pydantic ==2.13.*"
+
     - name: Optionally install OpenMM from PyPI
       if: ${{ matrix.install-source == 'pypi' && matrix.openmm == 'true' }}
       run: pip install "openmm ~=8.3"

docs/releasehistory.md

@@ -11,13 +11,19 @@ Please note that all releases prior to a version 1.0.0 are considered pre-releas
 
 ## 0.4.0
 
+### Behavior changes
+
+* #148 Switches version handling from `versioningit` to `setuptools-scm`
+* #174 Drops support for Python 3.11
+* #178 Define Pydantic schema on `Quantity`
+
+# Other improvements
+
 * #137 Add some NMR-related units
 * #138 Support PEP 639
 * #143 Updates how some data files are packaged
-* #148 Switches version handling from `versioningit` to `setuptools-scm`
 * #149 Runs tests with Python 3.14
 * #150 Runs tests with Pint 0.25
-* #174 Drops support for Python 3.11
 
 ## 0.3.1
 

openff/units/_tests/test_pydantic.py

@@ -0,0 +1,97 @@
+import numpy
+import pytest
+
+from openff.units import Quantity, unit
+
+pydantic = pytest.importorskip("pydantic")
+
+
+def test_model_definition():
+    """Just define a Pydantic model, which will crash if the schema is bad."""
+
+    class MyModel(pydantic.BaseModel):
+        x: Quantity
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        1.0 * unit.angstrom,
+        Quantity("1.0 angstrom"),
+        "1.0 angstrom",
+        {"magnitude": 1.0, "units": "angstrom"},
+    ],
+)
+def test_basic_field_validation(value):
+    class MyModel(pydantic.BaseModel):
+        x: Quantity
+
+    stored_value = MyModel(x=value).x
+
+    assert str(stored_value.units) == "angstrom"
+    assert stored_value.m_as("angstrom") == 1.0
+
+
+@pytest.mark.parametrize("serialize_with", ["python", "json"])
+def test_model_roundtrip(serialize_with):
+    """Test that a model can be round-tripped in Python."""
+
+    class MyModel(pydantic.BaseModel):
+        x: Quantity
+        y: Quantity
+        z: Quantity
+
+    model = MyModel(
+        x=1.0 * unit.angstrom,
+        y=[2, 3] * unit.amu,
+        z={"magnitude": 299.99, "units": "kelvin"},
+    )
+
+    match serialize_with:
+        case "json":
+            model_as_json = model.model_dump_json()
+            new_model = MyModel.model_validate_json(model_as_json)
+        case "python":
+            model_as_dict = model.model_dump()
+            new_model = MyModel.model_validate(model_as_dict)
+
+    assert new_model.x == Quantity(1.0, unit.angstrom), model.model_dump_json()
+    assert (new_model.y == Quantity([2, 3], unit.amu)).all()
+    assert type(new_model.y.m) is not list
+    assert type(new_model.y.m) is numpy.ndarray
+    assert new_model.z == Quantity(299.99, unit.kelvin)
+
+
+@pytest.mark.parametrize("serialize_with", ["python", "json"])
+def test_different_iterables(serialize_with):
+    class MyModel(pydantic.BaseModel):
+        a: Quantity
+        b: Quantity
+        c: Quantity
+
+    model = MyModel(
+        a=(300.0, 300.1) * unit.kelvin,
+        b=[12.011, 1.008, 1.008, 0.0] * unit.amu,
+        c=4.0 * numpy.eye(3) * unit.angstrom,
+    )
+
+    match serialize_with:
+        case "json":
+            model_as_json = model.model_dump_json()
+            new_model = MyModel.model_validate_json(model_as_json)
+        case "python":
+            model_as_dict = model.model_dump()
+            new_model = MyModel.model_validate(model_as_dict)
+
+    assert new_model.a.shape == (2,)
+    assert new_model.b.shape == (4,)
+    assert new_model.c.shape == (3, 3)
+
+    assert str(new_model.a.units) == "kelvin"
+    assert str(new_model.b.units) == "unified_atomic_mass_unit"
+    assert str(new_model.c.units) == "angstrom"
+
+    assert new_model.a[1].m == 300.1
+    assert new_model.b[0].m == 12.011
+    assert new_model.b[-1].m == 0.0
+    assert new_model.c[2, 2].m == 4.0

openff/units/units.py

@@ -2,21 +2,34 @@
 Core classes for OpenFF Units
 """
 
+from __future__ import annotations
+
+import json
 import uuid
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+import numpy  # possible to make this optional?
 import pint
 from openff.utilities import requires_package
 from pint import Measurement as _Measurement
 from pint import Quantity as _Quantity
 from pint import Unit as _Unit
+from pint.facets.plain.quantity import PlainQuantity as PintQuantity
 
 from openff.units.utilities import get_defaults_path
 
 if TYPE_CHECKING:
     import openmm.unit
 
+try:
+    from pydantic import GetCoreSchemaHandler
+    from pydantic_core import core_schema
+
+    has_pydantic = True
+except ImportError:
+    has_pydantic = False
+
 __all__ = (
     "DEFAULT_UNIT_REGISTRY",
     "Measurement",
@@ -32,7 +45,120 @@ class Unit(pint.UnitRegistry.Unit):
     pass
 
 
-class Quantity(pint.UnitRegistry.Quantity):
+if has_pydantic:
+
+    class _QuantityMixin:
+        @classmethod
+        def serialize(
+            cls,
+            v: PintQuantity,
+            info: core_schema.SerializationInfo | None = None,
+        ) -> dict | str | PintQuantity:
+            to_json = info is not None and info.mode_is_json()
+
+            if to_json:
+                magnitude = v.magnitude
+
+                # storing numpy arrays natively works fine in memory in Python,
+                # but must be list-ified when serializing to JSON.
+                if isinstance(magnitude, numpy.ndarray):
+                    magnitude = v.magnitude.tolist()
+
+                # TODO: I think this is necessary for handling unit-wrapped arrays, but it is
+                #       not so performant. Scalar quantities can be directly serialized to much
+                #       shorter strings
+                return json.dumps(
+                    {
+                        "magnitude": magnitude,
+                        "units": str(v.units),
+                    }
+                )
+
+            return {
+                "magnitude": v.magnitude,
+                "units": str(v.units),
+            }
+
+        @classmethod
+        def validate(
+            cls,
+            v: dict | str | PintQuantity,
+        ):
+            if isinstance(v, Quantity):
+                return v
+            elif isinstance(v, str):
+                # TODO: A significant wart is that we have to try to guess whether the string is a
+                #       JSON-serialized quantity or a simple string representation of a quantity.
+                #       For example: input of "0.9 nanometer" can be passed directly to the
+                #       Quantity constructor, but "{"magnitude": 0.9, "units": "nanometer"}" cannot
+                #       as it needs to be unwrapped. A better solution would require a better way
+                #       of serializing unit-wrapped arrays to JSON
+                if "{" in v:
+                    deserialized = json.loads(v)
+                    return Quantity(
+                        deserialized["magnitude"],
+                        deserialized["units"],
+                    )
+                else:
+                    return Quantity(v)
+            elif isinstance(v, dict):
+                return Quantity(v["magnitude"], v["units"])
+            else:
+                # this cannot be accessed with the current core_schema definition - the types of
+                # the `v` argument to this method **happen** to be identical to the supported types
+                # in the core_schema. If **either** is changed, this clause may be hit
+                raise ValueError(f"Invalid type {type(v)} for Quantity")
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            cls,
+            source_type: Any,
+            handler: GetCoreSchemaHandler,
+        ) -> core_schema.CoreSchema:
+
+            validate_schema = core_schema.chain_schema(
+                [
+                    core_schema.union_schema(
+                        [
+                            core_schema.is_instance_schema(PintQuantity),
+                            core_schema.str_schema(),
+                            core_schema.dict_schema(),
+                            # any other types that could be accepted by a quantity field?
+                        ]
+                    ),
+                    core_schema.no_info_plain_validator_function(cls.validate),
+                ]
+            )
+
+            validate_json_schema = core_schema.chain_schema(
+                [
+                    core_schema.union_schema(
+                        [
+                            core_schema.str_schema(coerce_numbers_to_str=True),
+                            core_schema.dict_schema(),
+                        ]
+                    ),
+                    core_schema.no_info_plain_validator_function(cls.validate),
+                ]
+            )
+
+            serialize_schema = core_schema.plain_serializer_function_ser_schema(
+                cls.serialize,
+                info_arg=True,
+            )
+
+            return core_schema.json_or_python_schema(
+                json_schema=validate_json_schema,
+                python_schema=validate_schema,
+                serialization=serialize_schema,
+            )
+else:
+
+    class _QuantityMixin:
+        pass
+
+
+class Quantity(_QuantityMixin, PintQuantity):
     """A value with associated units."""
 
     def __dask_tokenize__(self):
@@ -45,7 +171,7 @@ def _dask_finalize(results, func, args, units):
 
 
 @requires_package("openmm")
-def _to_openmm(self) -> "openmm.unit.Quantity":
+def _to_openmm(self) -> openmm.unit.Quantity:
     """Convert the quantity to an ``openmm.unit.Quantity``.
 
     Returns
@@ -84,10 +210,18 @@ class UnitRegistry(pint.UnitRegistry):
 Quantity: type[_Quantity] = DEFAULT_UNIT_REGISTRY.Quantity
 Measurement: type[_Measurement] = DEFAULT_UNIT_REGISTRY.Measurement
 
-pint.set_application_registry(DEFAULT_UNIT_REGISTRY)
-
 Quantity.to_openmm = _to_openmm  # type: ignore[attr-defined]
 
+if has_pydantic:
+    # Re-attach the Pydantic magic to our new Quantity class, which itself was
+    # dynamically created by Pint's magic (and lost these methods in the process).
+    Quantity.__get_pydantic_core_schema__ = _QuantityMixin.__get_pydantic_core_schema__  # type: ignore[attr-defined]
+    Quantity.validate = _QuantityMixin.validate  # type: ignore[attr-defined]
+    Quantity.serialize = _QuantityMixin.serialize  # type: ignore[attr-defined]
+
+
+pint.set_application_registry(DEFAULT_UNIT_REGISTRY)
+
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     Quantity([])