Source code for botorch.utils.dispatcher

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

from inspect import getsource, getsourcefile
from typing import Any, Callable, Optional, Tuple, Type

from multipledispatch.dispatcher import (
    Dispatcher as MDDispatcher,
    MDNotImplementedError,  # trivial subclass of NotImplementedError
    str_signature,
)


[docs] def type_bypassing_encoder(arg: Any) -> Type: # Allow type variables to be passed as pre-encoded arguments return arg if isinstance(arg, type) else type(arg)
[docs] class Dispatcher(MDDispatcher): r"""Clearing house for multiple dispatch functionality. This class extends `<multipledispatch.Dispatcher>` by: (i) generalizing the argument encoding convention during method lookup, (ii) implementing `__getitem__` as a dedicated method lookup function. """ def __init__( self, name: str, doc: Optional[str] = None, encoder: Callable[Any, Type] = type, ) -> None: """ Args: name: A string identifier for the `Dispatcher` instance. doc: A docstring for the multiply dispatched method(s). encoder: A callable that individually transforms the arguments passed at runtime in order to construct the key used for method lookup as `tuple(map(encoder, args))`. Defaults to `type`. """ super().__init__(name=name, doc=doc) self._encoder = encoder def __getitem__( self, args: Optional[Any] = None, types: Optional[Tuple[Type]] = None, ) -> Callable: r"""Method lookup. Args: args: A set of arguments that act as identifiers for a stored method. types: A tuple of types that encodes `args`. Returns: A callable corresponding to the given `args` or `types`. """ if types is None: if args is None: raise RuntimeError("One of `args` or `types` must be provided.") types = self.encode_args(args) elif args is not None: raise RuntimeError("Only one of `args` or `types` may be provided.") try: func = self._cache[types] except KeyError: func = self.dispatch(*types) if not func: msg = f"{self.name}: <{', '.join(cls.__name__ for cls in types)}" raise NotImplementedError(f"Could not find signature for {msg}") self._cache[types] = func return func def __call__(self, *args: Any, **kwargs: Any) -> Any: r"""Multiply dispatches a call to a collection of methods. Args: args: A set of arguments that act as identifiers for a stored method. kwargs: Optional keyword arguments passed to the retrieved method. Returns: The result of evaluating `func(*args, **kwargs)`, where `func` is the function obtained via method lookup. """ types = self.encode_args(args) func = self.__getitem__(types=types) try: return func(*args, **kwargs) except MDNotImplementedError: # Traverses registered methods in order, yields whenever a match is found funcs = self.dispatch_iter(*types) next(funcs) # burn first, same as self.__getitem__(types=types) for func in funcs: try: return func(*args, **kwargs) except MDNotImplementedError: pass raise NotImplementedError( f"Matching functions for {self.name:s}: {str_signature(types):s} " "found, but none completed successfully" )
[docs] def dispatch(self, *types: Type) -> Callable: r"""Method lookup strategy. Checks for an exact match before traversing the set of registered methods according to the current ordering. Args: types: A tuple of types that gets compared with the signatures of registered methods to determine compatibility. Returns: The first method encountered with a matching signature. """ if types in self.funcs: return self.funcs[types] try: return next(self.dispatch_iter(*types)) except StopIteration: return None
[docs] def encode_args(self, args: Any) -> Tuple[Type]: r"""Converts arguments into a tuple of types used during method lookup.""" return tuple(map(self.encoder, args if isinstance(args, tuple) else (args,)))
def _help(self, *args: Any) -> str: r"""Returns the retrieved method's docstring.""" return self.dispatch(*self.encode_args(args)).__doc__
[docs] def help(self, *args: Any, **kwargs: Any) -> None: r"""Prints the retrieved method's docstring.""" print(self._help(*args))
def _source(self, *args: Any) -> str: r"""Returns the retrieved method's source types as a string.""" func = self.dispatch(*self.encode_args(args)) if not func: raise TypeError("No function found") return f"File: {getsourcefile(func)}\n\n{getsource(func)}"
[docs] def source(self, *args, **kwargs) -> None: r"""Prints the retrieved method's source types.""" print(self._source(*args))
@property def encoder(self) -> Callable[Any, Type]: return self._encoder