De-stringify ParameterAttribute.unit (#1890)

* Reproduce #1888 * De-stringify ParameterAttribute.unit * Rename plugin class * Update release history
openforcefield · Jun 14, 2024 · fd5a20b · fd5a20b
1 parent f74700a
commit fd5a20b
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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",
         ]
     },
 )