From fd5a20b6094ac0f64f66510f2791244e9e5ef9b4 Mon Sep 17 00:00:00 2001
From: Matt Thompson <mattwthompson@protonmail.com>
Date: Fri, 14 Jun 2024 11:59:14 -0500
Subject: [PATCH] De-stringify `ParameterAttribute.unit` (#1890)

* Reproduce #1888

* De-stringify ParameterAttribute.unit

* Rename plugin class

* Update release history
---
 docs/releasehistory.md                        |  6 ++++
 .../toolkit/_tests/test_parameter_plugins.py  | 29 +++++++++++++++++--
 .../typing/engines/smirnoff/parameters.py     |  6 +++-
 .../custom_plugins/handler_plugins.py         | 29 ++++++++++++++++++-
 utilities/test_plugins/setup.py               |  1 +
 5 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/docs/releasehistory.md b/docs/releasehistory.md
index 046d03ad4..69083d6f3 100644
--- a/docs/releasehistory.md
+++ b/docs/releasehistory.md
@@ -10,6 +10,12 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
 * #1861 Forbid looking up virtual site parameters by SMIRKS, since a valid SMIRNOFF force field 
   can have several virtual site parameters with identical SMIRKS.
 
+## Current development
+
+### Bugfixes
+
+- [PR #1890](https://github.com/openforcefield/openff-toolkit/pull/1890): Fixes [#1888](https://github.com/openforcefield/openff-toolkit/pull/1888) in which `ParameterAttribute.unit` was sometimes not validated when stringified.
+
 ## 0.16.1
 
 
diff --git a/openff/toolkit/_tests/test_parameter_plugins.py b/openff/toolkit/_tests/test_parameter_plugins.py
index e860ffcda..a27734098 100644
--- a/openff/toolkit/_tests/test_parameter_plugins.py
+++ b/openff/toolkit/_tests/test_parameter_plugins.py
@@ -4,7 +4,7 @@
 
 import pytest
 
-from openff.toolkit.typing.engines.smirnoff import ForceField
+from openff.toolkit import ForceField, Quantity
 from openff.toolkit.typing.engines.smirnoff.plugins import (
     _load_handler_plugins,
     load_handler_plugins,
@@ -50,8 +50,12 @@ def test_load_handler_plugins():
 
     registered_plugins = load_handler_plugins()
 
-    assert len(registered_plugins) == 1
-    assert registered_plugins[0].__name__ == "CustomHandler"
+    assert len(registered_plugins) == 2
+
+    registered_plugin_names = [plugin.__name__ for plugin in registered_plugins]
+
+    assert "CustomHandler" in registered_plugin_names
+    assert "FOOBuckinghamHandler" in registered_plugin_names
 
 
 def test_do_not_load_other_type():
@@ -67,3 +71,22 @@ def test_skip_wrong_subclass(caplog):
     load_handler_plugins()
 
     assert "does not inherit from ParameterHandler" in caplog.text, caplog.text
+
+
+def test_buckingham_type():
+    """Reproduce, in part, issue #1888."""
+    from custom_plugins.handler_plugins import FOOBuckinghamHandler
+
+    parameter = FOOBuckinghamHandler.FOOBuckinghamType(
+        smirks="[*:1]",
+        a="2 kilojoule_per_mole",
+        b="1/nanometer",
+        c="-0.5 kilojoule_per_mole * nanometer**6",
+    )
+
+    for param in ["a", "b", "c"]:
+        assert isinstance(getattr(parameter, param), Quantity)
+
+    assert str(parameter.a.units) == "kilojoule_per_mole"
+    assert str(parameter.b.units) == "1 / nanometer"
+    assert str(parameter.c.units) == "kilojoule_per_mole * nanometer ** 6"
diff --git a/openff/toolkit/typing/engines/smirnoff/parameters.py b/openff/toolkit/typing/engines/smirnoff/parameters.py
index caa082ff2..afbc6f1aa 100644
--- a/openff/toolkit/typing/engines/smirnoff/parameters.py
+++ b/openff/toolkit/typing/engines/smirnoff/parameters.py
@@ -342,10 +342,14 @@ class ParameterAttribute:
     def __init__(
         self,
         default: Any = UNDEFINED,
-        unit: Optional[Unit] = None,
+        unit: Unit | str | None = None,
         converter: Optional[Callable] = None,
         docstring: str = "",
     ):
+        if isinstance(unit, str):
+            # be careful with module & variable names
+            unit = Unit(unit)
+
         self.default = default
         self._unit = unit
         self._converter = converter
diff --git a/utilities/test_plugins/custom_plugins/handler_plugins.py b/utilities/test_plugins/custom_plugins/handler_plugins.py
index af239ba12..ff0fc4646 100644
--- a/utilities/test_plugins/custom_plugins/handler_plugins.py
+++ b/utilities/test_plugins/custom_plugins/handler_plugins.py
@@ -1,5 +1,6 @@
 from openff.toolkit.typing.engines.smirnoff import ParameterHandler, ParameterIOHandler
-
+from openff.toolkit.typing.engines.smirnoff.parameters import ParameterAttribute, ParameterType, _NonbondedHandler
+from openff.toolkit import unit
 
 class CustomHandler(ParameterHandler):
     _TAGNAME = "CustomHandler"
@@ -11,3 +12,29 @@ class WrongSubclass(list):
 
 class CustomIOHandler(ParameterIOHandler):
     _FORMAT = "JSON"
+
+class FOOBuckinghamHandler(_NonbondedHandler):
+    """A custom parameter handler for buckingham interactions."""
+
+    class FOOBuckinghamType(ParameterType):
+        """A custom parameter type for buckingham interactions."""
+
+        _ELEMENT_NAME = "Atom"
+
+        # Define unit as a Unit object
+        a = ParameterAttribute(default=None,
+                               unit=unit.kilojoule_per_mole,
+            )
+
+        # Define using a string
+        b = ParameterAttribute(default=None,
+                               unit="nanometer**-1",
+                               )
+
+        c = ParameterAttribute(
+            default=None,
+            unit=unit.kilojoule_per_mole * unit.nanometer ** 6,
+        )
+
+    _TAGNAME = "FOOBuckingham"
+    _INFOTYPE = FOOBuckinghamType
diff --git a/utilities/test_plugins/setup.py b/utilities/test_plugins/setup.py
index b5846717a..cc6dca6d6 100644
--- a/utilities/test_plugins/setup.py
+++ b/utilities/test_plugins/setup.py
@@ -12,6 +12,7 @@
         "openff.toolkit.plugins.handlers": [
             "CustomHandler = custom_plugins.handler_plugins:CustomHandler",
             "WrongSubclass = custom_plugins.handler_plugins:WrongSubclass",
+            "FOOBuckinghamHandler = custom_plugins.handler_plugins:FOOBuckinghamHandler",
         ]
     },
 )