Un-vendor toolkit_registry_manager

openfe/utils/remove_oechem.py

@@ -1,28 +1,23 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
-from openff.toolkit import GLOBAL_TOOLKIT_REGISTRY, OpenEyeToolkitWrapper
-from openff.toolkit.utils.toolkit_registry import ToolkitUnavailableException
+from openff.toolkit.utils.toolkit_registry import (
+    ToolkitRegistry,
+    toolkit_registry_manager,
+)
+from openff.toolkit.utils.toolkits import (
+    AmberToolsToolkitWrapper,
+    BuiltInToolkitWrapper,
+    RDKitToolkitWrapper,
+)
 
 from contextlib import contextmanager
 
-
-@contextmanager
-def without_oechem_backend():
-    """For temporarily removing oechem from openff's toolkit registry"""
-    current_toolkits = [type(tk)
-                        for tk in GLOBAL_TOOLKIT_REGISTRY.registered_toolkits]
-
-    try:
-        GLOBAL_TOOLKIT_REGISTRY.deregister_toolkit(OpenEyeToolkitWrapper())
-    except ToolkitUnavailableException:
-        pass
-
-    try:
-        yield None
-    finally:
-        # this is order dependent; we want to prepend OEChem back to first
-        while GLOBAL_TOOLKIT_REGISTRY.registered_toolkits:
-            GLOBAL_TOOLKIT_REGISTRY.deregister_toolkit(
-                GLOBAL_TOOLKIT_REGISTRY.registered_toolkits[0])
-        for tk in current_toolkits:
-            GLOBAL_TOOLKIT_REGISTRY.register_toolkit(tk)
+without_oechem_backend = toolkit_registry_manager(
+    toolkit_registry=ToolkitRegistry(
+        toolkit_precedence=[
+            RDKitToolkitWrapper(),
+            AmberToolsToolkitWrapper(),
+            BuiltInToolkitWrapper(),
+        ]
+    )
+)