From ed31fc343e2b17ba7ace1d9e32fe2cc4225ea9b1 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Thu, 8 Jun 2023 18:02:38 -0500
Subject: [PATCH 1/2] MAINT: Test on multiple Pint versions

---
 .github/workflows/ci.yaml         | 2 ++
 devtools/conda-envs/test_env.yaml | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index adb6ca5..b629235 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -22,6 +22,7 @@ jobs:
         os: [macOS-latest, ubuntu-latest]
         openmm: ["true", "false"]
         python-version: ["3.9", "3.10", "3.11"]
+        pint-version: ["0.21", "0.22"]
 
     env:
       CI_OS: ${{ matrix.os }}
@@ -38,6 +39,7 @@ jobs:
         environment-file: devtools/conda-envs/test_env.yaml
         create-args: >-
           python=${{ matrix.python-version }}
+          pint=${{ matrix.pint-version }}
 
     - name: Optionally install OpenMM
       if: ${{ matrix.openmm == 'true' }}
diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml
index 1d85b51..8602cb6 100644
--- a/devtools/conda-envs/test_env.yaml
+++ b/devtools/conda-envs/test_env.yaml
@@ -6,7 +6,7 @@ dependencies:
   - python
   - pip
   - numpy
-  - pint =0.21  # https://github.com/conda-forge/pint-feedstock/pull/57
+  - pint
   - openff-utilities >=0.1.3
 
     # Tests

From 8f655d6f64a309fe26b0390571033d9c2a2778d8 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Wed, 14 Jun 2023 15:31:40 -0500
Subject: [PATCH 2/2] FIX: Update monkey-patch, typing

---
 openff/units/elements.py |  3 ++-
 openff/units/openmm.py   |  3 ++-
 openff/units/units.py    | 34 ++++++++++++++++++----------------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/openff/units/elements.py b/openff/units/elements.py
index 29b55b5..c88cef3 100644
--- a/openff/units/elements.py
+++ b/openff/units/elements.py
@@ -37,7 +37,8 @@
 ]
 
 MASSES: Dict[int, Quantity] = {
-    index + 1: Quantity(mass, unit.dalton)
+    # https://github.com/hgrecco/pint/issues/1804
+    index + 1: Quantity(mass, unit.dalton)  # type: ignore[call-overload]
     for index, mass in enumerate(
         [
             1.007947,
diff --git a/openff/units/openmm.py b/openff/units/openmm.py
index b83b89a..2880758 100644
--- a/openff/units/openmm.py
+++ b/openff/units/openmm.py
@@ -264,7 +264,8 @@ def _ensure_openff_quantity(
             )
     else:
         try:
-            return unit.Quantity(
+            # https://github.com/hgrecco/pint/issues/1804
+            return unit.Quantity(  # type: ignore[call-overload]
                 unknown_quantity,
                 unit.dimensionless,
             )
diff --git a/openff/units/units.py b/openff/units/units.py
index 4ca257d..9858baf 100644
--- a/openff/units/units.py
+++ b/openff/units/units.py
@@ -10,7 +10,6 @@
 from pint import Measurement as _Measurement
 from pint import Quantity as _Quantity
 from pint import Unit as _Unit
-from pint import UnitRegistry as _UnitRegistry
 
 from openff.units.utilities import get_defaults_path
 
@@ -25,13 +24,13 @@
 ]
 
 
-class Unit(_Unit):
+class Unit(pint.UnitRegistry.Unit):
     """A unit of measure."""
 
     pass
 
 
-class Quantity(_Quantity):
+class Quantity(pint.UnitRegistry.Quantity):
     """A value with associated units."""
 
     def __dask_tokenize__(self):
@@ -42,21 +41,22 @@ def _dask_finalize(results, func, args, units):
         values = func(results, *args)
         return Quantity(values, units)
 
-    @requires_package("openmm")
-    def to_openmm(self) -> "OpenMMQuantity":
-        """Convert the quantity to an ``openmm.unit.Quantity``.
 
-        Returns
-        -------
-        openmm_quantity : openmm.unit.quantity.Quantity
-            The OpenMM compatible quantity.
-        """
-        from openff.units.openmm import to_openmm
+@requires_package("openmm")
+def _to_openmm(self) -> "OpenMMQuantity":
+    """Convert the quantity to an ``openmm.unit.Quantity``.
 
-        return to_openmm(self)
+    Returns
+    -------
+    openmm_quantity : openmm.unit.quantity.Quantity
+        The OpenMM compatible quantity.
+    """
+    from openff.units.openmm import to_openmm
 
+    return to_openmm(self)
 
-class Measurement(_Measurement):
+
+class Measurement(pint.UnitRegistry.Measurement):  # type: ignore
     """A value with associated units and uncertainty."""
 
     def __dask_tokenize__(self):
@@ -68,7 +68,7 @@ def _dask_finalize(results, func, args, units):
         return Measurement(values, units)
 
 
-class UnitRegistry(_UnitRegistry):
+class UnitRegistry(pint.UnitRegistry):
     _quantity_class = Quantity
     _unit_class = Unit
     _measurement_class = Measurement
@@ -78,10 +78,12 @@ class UnitRegistry(_UnitRegistry):
 
 Unit: _Unit = DEFAULT_UNIT_REGISTRY.Unit  # type: ignore[no-redef]
 Quantity: _Quantity = DEFAULT_UNIT_REGISTRY.Quantity  # type: ignore[no-redef]
-Measurement: _Measurement = DEFAULT_UNIT_REGISTRY.Measurement  # type: ignore[no-redef]
+Measurement: _Measurement = DEFAULT_UNIT_REGISTRY.Measurement  # type: ignore
 
 pint.set_application_registry(DEFAULT_UNIT_REGISTRY)
 
+Quantity.to_openmm = _to_openmm  # type: ignore[attr-defined]
+
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     Quantity([])