From 83dc8b45b6616a94c60a6c1c87fcfb51d5027fb3 Mon Sep 17 00:00:00 2001
From: "Matthew W. Thompson" <mattwthompson@protonmail.com>
Date: Tue, 13 Feb 2024 11:43:49 -0600
Subject: [PATCH] Un-vendor `toolkit_registry_manager`

---
 openfe/utils/remove_oechem.py | 41 +++++++++++++++--------------------
 1 file changed, 18 insertions(+), 23 deletions(-)

diff --git a/openfe/utils/remove_oechem.py b/openfe/utils/remove_oechem.py
index f489f04cc..38c019cbd 100644
--- a/openfe/utils/remove_oechem.py
+++ b/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 = toolkit_registry_manager(
+    toolkit_registry=ToolkitRegistry(
+        toolkit_precedence=[
+            RDKitToolkitWrapper(),
+            AmberToolsToolkitWrapper(),
+            BuiltInToolkitWrapper(),
+        ]
+    )
+)