Source code for openff.toolkit.utils.toolkit_registry
"Registry for ToolkitWrapper objects"
__all__ = ("ToolkitRegistry", "toolkit_registry_manager")
import inspect
import logging
from contextlib import contextmanager
from typing import Callable, Optional, Union
from openff.toolkit.utils.ambertools_wrapper import AmberToolsToolkitWrapper
from openff.toolkit.utils.base_wrapper import ToolkitWrapper
from openff.toolkit.utils.builtin_wrapper import BuiltInToolkitWrapper
from openff.toolkit.utils.exceptions import (
InvalidToolkitError,
LicenseError,
ToolkitUnavailableException,
)
from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper
from openff.toolkit.utils.rdkit_wrapper import RDKitToolkitWrapper
from openff.toolkit.utils.utils import all_subclasses
logger = logging.getLogger(__name__)
[docs]class ToolkitRegistry:
"""
Registry for ToolkitWrapper objects
Examples
--------
Register toolkits in a specified order, skipping if unavailable
>>> from openff.toolkit.utils.toolkits import ToolkitRegistry
>>> toolkit_precedence = [OpenEyeToolkitWrapper, RDKitToolkitWrapper, AmberToolsToolkitWrapper]
>>> toolkit_registry = ToolkitRegistry(toolkit_precedence)
>>> toolkit_registry
<ToolkitRegistry containing OpenEye Toolkit, The RDKit, AmberTools>
Register all available toolkits (in the order OpenEye, RDKit, AmberTools, built-in)
>>> toolkits = [OpenEyeToolkitWrapper, RDKitToolkitWrapper, AmberToolsToolkitWrapper, BuiltInToolkitWrapper]
>>> toolkit_registry = ToolkitRegistry(toolkit_precedence=toolkits)
>>> toolkit_registry
<ToolkitRegistry containing OpenEye Toolkit, The RDKit, AmberTools, Built-in Toolkit>
Retrieve the global singleton toolkit registry, which is created when this module is imported from all available
toolkits:
>>> from openff.toolkit.utils.toolkits import GLOBAL_TOOLKIT_REGISTRY as toolkit_registry
>>> toolkit_registry
<ToolkitRegistry containing OpenEye Toolkit, The RDKit, AmberTools, Built-in Toolkit>
Note that this will contain different ToolkitWrapper objects based on what toolkits
are currently installed.
.. warning :: This API is experimental and subject to change.
"""
_toolkits: list[ToolkitWrapper]
[docs] def __init__(
self,
toolkit_precedence: Optional[list[type[ToolkitWrapper]]] = None,
exception_if_unavailable: bool = True,
_register_imported_toolkit_wrappers: bool = False,
):
"""
Create an empty toolkit registry.
Parameters
----------
toolkit_precedence
List of toolkit wrapper classes, in order of desired precedence when performing molecule operations. If
None, no toolkits will be registered.
exception_if_unavailable
If True, an exception will be raised if the toolkit is unavailable
_register_imported_toolkit_wrappers
If True, will attempt to register all imported ToolkitWrapper subclasses that can be
found in the order of toolkit_precedence, if specified. If toolkit_precedence is not
specified, the default order is [OpenEyeToolkitWrapper, RDKitToolkitWrapper,
AmberToolsToolkitWrapper, BuiltInToolkitWrapper].
"""
self._toolkits = list()
toolkit_classes_to_register: list[type[ToolkitWrapper]] = list()
if _register_imported_toolkit_wrappers:
if toolkit_precedence is None:
toolkit_precedence = [
OpenEyeToolkitWrapper,
RDKitToolkitWrapper,
AmberToolsToolkitWrapper,
BuiltInToolkitWrapper,
]
all_importable_toolkit_wrappers = all_subclasses(ToolkitWrapper)
for toolkit in toolkit_precedence:
if toolkit in all_importable_toolkit_wrappers:
toolkit_classes_to_register.append(toolkit)
else:
if toolkit_precedence is not None:
toolkit_classes_to_register = toolkit_precedence
if toolkit_classes_to_register:
for toolkit in toolkit_classes_to_register:
self.register_toolkit(
toolkit_wrapper=toolkit,
exception_if_unavailable=exception_if_unavailable,
)
@property
def registered_toolkits(self) -> list[ToolkitWrapper]:
"""
List registered toolkits.
.. warning :: This API is experimental and subject to change.
.. todo :: Should this return a generator? Deep copies? Classes? Toolkit names?
Returns
-------
toolkits
"""
return list(self._toolkits)
@property
def registered_toolkit_versions(self) -> dict[str, str]:
"""
Return a dict containing the version of each registered toolkit.
.. warning :: This API is experimental and subject to change.
Returns
-------
toolkit_versions
A dictionary mapping names and versions of wrapped toolkits
"""
return {tk.toolkit_name: tk.toolkit_version for tk in self.registered_toolkits}
[docs] def register_toolkit(
self,
toolkit_wrapper: Union[ToolkitWrapper, type[ToolkitWrapper]],
exception_if_unavailable: bool = True,
):
"""
Register the provided toolkit wrapper class, instantiating an object of it.
.. warning :: This API is experimental and subject to change.
.. todo ::
This method should raise an exception if the toolkit is unavailable, unless an optional argument
is specified that silently avoids registration of toolkits that are unavailable.
Parameters
----------
toolkit_wrapper
The toolkit wrapper to register or its class.
exception_if_unavailable
If True, an exception will be raised if the toolkit is unavailable
"""
# Instantiate class if class, or just add if already instantiated.
if isinstance(toolkit_wrapper, ToolkitWrapper):
self._toolkits.append(toolkit_wrapper)
elif issubclass(toolkit_wrapper, ToolkitWrapper):
try:
_toolkit_wrapper = toolkit_wrapper()
# This exception can be raised by OpenEyeToolkitWrapper
except LicenseError as license_exception:
if exception_if_unavailable:
raise ToolkitUnavailableException(license_exception.msg)
else:
logger.warning(license_exception)
return
except ToolkitUnavailableException:
if exception_if_unavailable:
raise ToolkitUnavailableException(
"Unable to load toolkit '{_toolkit_wrapper._toolkit_name}'. "
)
return
self._toolkits.append(_toolkit_wrapper)
else:
raise ValueError(f"Given unexpected argument {type(toolkit_wrapper)=}")
[docs] def deregister_toolkit(self, toolkit_wrapper: ToolkitWrapper):
"""
Remove a ToolkitWrapper from the list of toolkits in this ToolkitRegistry
.. warning :: This API is experimental and subject to change.
Parameters
----------
toolkit_wrapper
The toolkit wrapper to remove from the registry
Raises
------
InvalidToolkitError
If toolkit_wrapper is not a ToolkitWrapper or subclass
ToolkitUnavailableException
If toolkit_wrapper is not found in the registry
"""
# If passed a class, instantiate it
if inspect.isclass(toolkit_wrapper):
toolkit_wrapper = toolkit_wrapper()
if not isinstance(toolkit_wrapper, ToolkitWrapper):
msg = (
f"Argument {toolkit_wrapper} must an ToolkitWrapper "
f"or subclass of it. Found type {type(toolkit_wrapper)}."
)
raise InvalidToolkitError(msg)
toolkits_to_remove = []
for toolkit in self._toolkits:
if type(toolkit) is type(toolkit_wrapper):
toolkits_to_remove.append(toolkit)
if not toolkits_to_remove:
msg = (
f"Did not find {toolkit_wrapper} in registry. "
f"Currently registered toolkits are {self._toolkits}"
)
raise ToolkitUnavailableException(msg)
for toolkit_to_remove in toolkits_to_remove:
self._toolkits.remove(toolkit_to_remove)
[docs] def add_toolkit(self, toolkit_wrapper: ToolkitWrapper):
"""
Append a ToolkitWrapper onto the list of toolkits in this ToolkitRegistry
.. warning :: This API is experimental and subject to change.
Parameters
----------
toolkit_wrapper
The ToolkitWrapper object to add to the list of registered toolkits
Raises
------
InvalidToolkitError
If toolkit_wrapper is not a ToolkitWrapper or subclass
"""
if not isinstance(toolkit_wrapper, ToolkitWrapper):
raise InvalidToolkitError(
"Something other than a ToolkitWrapper object was passed to ToolkitRegistry.add_toolkit()\n"
f"Given object {toolkit_wrapper} of type {type(toolkit_wrapper)}"
)
self._toolkits.append(toolkit_wrapper)
# TODO: Can we automatically resolve calls to methods that are not explicitly defined using some Python magic?
[docs] def resolve(self, method_name: str) -> Callable:
"""
Resolve the requested method name by checking all registered toolkits in
order of precedence for one that provides the requested method.
Parameters
----------
method_name
The name of the method to resolve
Returns
-------
method
The method of the first registered toolkit that provides the requested method name
Raises
------
NotImplementedError if the requested method cannot be found among the registered toolkits
Examples
--------
Create a molecule, and call the toolkit ``to_smiles()`` method directly
>>> from openff.toolkit import Molecule
>>> molecule = Molecule.from_smiles('Cc1ccccc1')
>>> toolkit_registry = ToolkitRegistry([OpenEyeToolkitWrapper, RDKitToolkitWrapper, AmberToolsToolkitWrapper])
>>> method = toolkit_registry.resolve('to_smiles')
>>> smiles = method(molecule)
.. todo :: Is there a better way to figure out which toolkits implement given methods by introspection?
"""
for toolkit in self._toolkits:
if hasattr(toolkit, method_name):
method = getattr(toolkit, method_name)
return method
# No toolkit was found to provide the requested capability
# TODO: Can we help developers by providing a check for typos in expected method names?
raise NotImplementedError(
f'No registered toolkits can provide the capability "{method_name}".\n'
f"Available toolkits are: {self.registered_toolkits}\n"
)
# TODO: Can we instead register available methods directly with `ToolkitRegistry`,
# so we can just use `ToolkitRegistry.method()`?
[docs] def call(
self,
method_name: str,
*args,
raise_exception_types: Optional[list[type[Exception]]] = None,
**kwargs,
):
"""
Execute the requested method by attempting to use all registered toolkits in order of precedence.
``*args`` and ``**kwargs`` are passed to the desired method, and return values of the method are returned
This is a convenient shorthand for ``toolkit_registry.resolve_method(method_name)(*args, **kwargs)``
Parameters
----------
method_name
The name of the method to execute
raise_exception_types
A list of exception-derived types to catch and raise immediately. If None, this will be set to [Exception],
which will raise an error immediately if the first ToolkitWrapper in the registry fails. To try each
ToolkitWrapper that provides a suitably-named method, set this to the empty list ([]). If all
ToolkitWrappers run without raising any exceptions in this list, a single ValueError will be raised
containing the each ToolkitWrapper that was tried and the exception it raised.
Raises
------
NotImplementedError if the requested method cannot be found among the registered toolkits
ValueError if no exceptions in the raise_exception_types list were raised by ToolkitWrappers, and
all ToolkitWrappers in the ToolkitRegistry were tried.
Other forms of exceptions are possible if raise_exception_types is specified.
These are defined by the ToolkitWrapper method being called.
Examples
--------
Create a molecule, and call the toolkit ``to_smiles()`` method directly
>>> from openff.toolkit import Molecule
>>> molecule = Molecule.from_smiles('Cc1ccccc1')
>>> toolkit_registry = ToolkitRegistry([OpenEyeToolkitWrapper, RDKitToolkitWrapper])
>>> smiles = toolkit_registry.call('to_smiles', molecule)
"""
if raise_exception_types is None:
raise_exception_types = [Exception]
errors = list()
for toolkit in self._toolkits:
if hasattr(toolkit, method_name):
method = getattr(toolkit, method_name)
try:
return method(*args, **kwargs)
except Exception as e:
for exception_type in raise_exception_types:
if isinstance(e, exception_type):
raise e
errors.append((toolkit, e))
# No toolkit was found to provide the requested capability
# TODO: Can we help developers by providing a check for typos in expected method names?
msg = (
f'No registered toolkits can provide the capability "{method_name}" '
f'for args "{args}" and kwargs "{kwargs}"\n'
)
msg += "Available toolkits are: {}\n".format(self.registered_toolkits)
# Append information about toolkits that implemented the method, but could not handle the provided parameters
for toolkit, error in errors:
msg += f" {toolkit} {type(error)} : {error}\n"
raise ValueError(msg)
def __repr__(self):
return f"<ToolkitRegistry containing {', '.join([tk.toolkit_name for tk in self._toolkits])}>"
# Copied from https://github.com/openforcefield/openff-fragmenter/blob/4a290b866a8ed43eabcbd3231c62b01f0c6d7df6
# /openff/fragmenter/utils.py#L97-L123
[docs]@contextmanager
def toolkit_registry_manager(toolkit_registry: Union[ToolkitRegistry, ToolkitWrapper]):
"""
A context manager that temporarily changes the ToolkitWrappers in the GLOBAL_TOOLKIT_REGISTRY.
This can be useful in cases where one would otherwise need to otherwise manually specify the use of a specific
ToolkitWrapper or ToolkitRegistry repeatedly in a block of code, or in cases where there isn't another way to
switch the ToolkitWrappers used for a particular operation.
Examples
--------
>>> from openff.toolkit import Molecule, RDKitToolkitWrapper, AmberToolsToolkitWrapper
>>> from openff.toolkit.utils import toolkit_registry_manager, ToolkitRegistry
>>> mol = Molecule.from_smiles("CCO")
>>> print(mol.to_smiles()) # This will use the OpenEye backend (if installed and licensed)
[H]C([H])([H])C([H])([H])O[H]
>>> with toolkit_registry_manager(ToolkitRegistry([RDKitToolkitWrapper()])):
... print(mol.to_smiles())
[H][O][C]([H])([H])[C]([H])([H])[H]
"""
from openff.toolkit.utils.toolkits import GLOBAL_TOOLKIT_REGISTRY
if isinstance(toolkit_registry, ToolkitRegistry):
context_toolkits = toolkit_registry.registered_toolkits
elif isinstance(toolkit_registry, ToolkitWrapper):
context_toolkits = [toolkit_registry]
else:
raise NotImplementedError(
"Only ``ToolkitRegistry`` and ``ToolkitWrapper`` are supported."
)
original_toolkits = GLOBAL_TOOLKIT_REGISTRY.registered_toolkits
for toolkit in original_toolkits:
GLOBAL_TOOLKIT_REGISTRY.deregister_toolkit(toolkit)
for toolkit in context_toolkits:
GLOBAL_TOOLKIT_REGISTRY.register_toolkit(toolkit)
try:
yield
finally:
for toolkit in context_toolkits:
GLOBAL_TOOLKIT_REGISTRY.deregister_toolkit(toolkit)
for toolkit in original_toolkits:
GLOBAL_TOOLKIT_REGISTRY.register_toolkit(toolkit)
_toolkit_registry_manager = toolkit_registry_manager