Source code for uarray._backend

import typing
import types
import inspect
import functools
from . import _uarray
import copyreg
import atexit
import pickle
import contextlib

ArgumentExtractorType = typing.Callable[..., typing.Tuple["Dispatchable", ...]]
ArgumentReplacerType = typing.Callable[
    [typing.Tuple, typing.Dict, typing.Tuple], typing.Tuple[typing.Tuple, typing.Dict]
]

from ._uarray import (  # type: ignore
    BackendNotImplementedError,
    _Function,
    _SkipBackendContext,
    _SetBackendContext,
    _BackendState,
)

__all__ = [
    "set_backend",
    "set_global_backend",
    "skip_backend",
    "register_backend",
    "determine_backend",
    "determine_backend_multi",
    "clear_backends",
    "create_multimethod",
    "generate_multimethod",
    "_Function",
    "BackendNotImplementedError",
    "Dispatchable",
    "wrap_single_convertor",
    "wrap_single_convertor_instance",
    "all_of_type",
    "mark_as",
    "set_state",
    "get_state",
    "reset_state",
    "_BackendState",
    "_SkipBackendContext",
    "_SetBackendContext",
]


def unpickle_function(mod_name, qname, self_):
    import importlib

    try:
        module = importlib.import_module(mod_name)
        qname = qname.split(".")
        func = module
        for q in qname:
            func = getattr(func, q)

        if self_ is not None:
            func = types.MethodType(func, self_)

        return func
    except (ImportError, AttributeError) as e:
        from pickle import UnpicklingError

        raise UnpicklingError from e


def pickle_function(func):
    mod_name = getattr(func, "__module__", None)
    qname = getattr(func, "__qualname__", None)
    self_ = getattr(func, "__self__", None)

    try:
        test = unpickle_function(mod_name, qname, self_)
    except pickle.UnpicklingError:
        test = None

    if test is not func:
        raise pickle.PicklingError(
            "Can't pickle {}: it's not the same object as {}".format(func, test)
        )

    return unpickle_function, (mod_name, qname, self_)


def pickle_state(state):
    return _uarray._BackendState._unpickle, state._pickle()


def pickle_set_backend_context(ctx):
    return _SetBackendContext, ctx._pickle()


def pickle_skip_backend_context(ctx):
    return _SkipBackendContext, ctx._pickle()


copyreg.pickle(_Function, pickle_function)
copyreg.pickle(_uarray._BackendState, pickle_state)
copyreg.pickle(_SetBackendContext, pickle_set_backend_context)
copyreg.pickle(_SkipBackendContext, pickle_skip_backend_context)
atexit.register(_uarray.clear_all_globals)


[docs]def get_state(): """ Returns an opaque object containing the current state of all the backends. Can be used for synchronization between threads/processes. See Also -------- set_state Sets the state returned by this function. """ return _uarray.get_state()
[docs]@contextlib.contextmanager def reset_state(): """ Returns a context manager that resets all state once exited. See Also -------- set_state Context manager that sets the backend state. get_state Gets a state to be set by this context manager. """ with set_state(get_state()): yield
[docs]@contextlib.contextmanager def set_state(state): """ A context manager that sets the state of the backends to one returned by :obj:`get_state`. See Also -------- get_state Gets a state to be set by this context manager. """ old_state = get_state() _uarray.set_state(state) try: yield finally: _uarray.set_state(old_state, True)
[docs]def create_multimethod(*args, **kwargs): """ Creates a decorator for generating multimethods. This function creates a decorator that can be used with an argument extractor in order to generate a multimethod. Other than for the argument extractor, all arguments are passed on to :obj:`generate_multimethod`. See Also -------- generate_multimethod Generates a multimethod. """ def wrapper(a): return generate_multimethod(a, *args, **kwargs) return wrapper
[docs]def generate_multimethod( argument_extractor: ArgumentExtractorType, argument_replacer: ArgumentReplacerType, domain: str, default: typing.Optional[typing.Callable] = None, ): """ Generates a multimethod. Parameters ---------- argument_extractor : ArgumentExtractorType A callable which extracts the dispatchable arguments. Extracted arguments should be marked by the :obj:`Dispatchable` class. It has the same signature as the desired multimethod. argument_replacer : ArgumentReplacerType A callable with the signature (args, kwargs, dispatchables), which should also return an (args, kwargs) pair with the dispatchables replaced inside the args/kwargs. domain : str A string value indicating the domain of this multimethod. default: Optional[Callable], optional The default implementation of this multimethod, where ``None`` (the default) specifies there is no default implementation. Examples -------- In this example, ``a`` is to be dispatched over, so we return it, while marking it as an ``int``. The trailing comma is needed because the args have to be returned as an iterable. >>> def override_me(a, b): ... return Dispatchable(a, int), Next, we define the argument replacer that replaces the dispatchables inside args/kwargs with the supplied ones. >>> def override_replacer(args, kwargs, dispatchables): ... return (dispatchables[0], args[1]), {} Next, we define the multimethod. >>> overridden_me = generate_multimethod( ... override_me, override_replacer, "ua_examples" ... ) Notice that there's no default implementation, unless you supply one. >>> overridden_me(1, "a") Traceback (most recent call last): ... uarray.BackendNotImplementedError: ... >>> overridden_me2 = generate_multimethod( ... override_me, override_replacer, "ua_examples", default=lambda x, y: (x, y) ... ) >>> overridden_me2(1, "a") (1, 'a') See Also -------- uarray See the module documentation for how to override the method by creating backends. """ kw_defaults, arg_defaults, opts = get_defaults(argument_extractor) ua_func = _Function( argument_extractor, argument_replacer, domain, arg_defaults, kw_defaults, default, ) return functools.update_wrapper(ua_func, argument_extractor)
[docs]def set_backend(backend, coerce=False, only=False): """ A context manager that sets the preferred backend. Parameters ---------- backend The backend to set. coerce Whether or not to coerce to a specific backend's types. Implies ``only``. only Whether or not this should be the last backend to try. See Also -------- skip_backend: A context manager that allows skipping of backends. set_global_backend: Set a single, global backend for a domain. """ try: return backend.__ua_cache__["set", coerce, only] except AttributeError: backend.__ua_cache__ = {} except KeyError: pass ctx = _SetBackendContext(backend, coerce, only) backend.__ua_cache__["set", coerce, only] = ctx return ctx
[docs]def skip_backend(backend): """ A context manager that allows one to skip a given backend from processing entirely. This allows one to use another backend's code in a library that is also a consumer of the same backend. Parameters ---------- backend The backend to skip. See Also -------- set_backend: A context manager that allows setting of backends. set_global_backend: Set a single, global backend for a domain. """ try: return backend.__ua_cache__["skip"] except AttributeError: backend.__ua_cache__ = {} except KeyError: pass ctx = _SkipBackendContext(backend) backend.__ua_cache__["skip"] = ctx return ctx
def get_defaults(f): sig = inspect.signature(f) kw_defaults = {} arg_defaults = [] opts = set() for k, v in sig.parameters.items(): if v.default is not inspect.Parameter.empty: kw_defaults[k] = v.default if v.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ): arg_defaults.append(v.default) opts.add(k) return kw_defaults, tuple(arg_defaults), opts
[docs]def set_global_backend(backend, coerce=False, only=False, *, try_last=False): """ This utility method replaces the default backend for permanent use. It will be tried in the list of backends automatically, unless the ``only`` flag is set on a backend. This will be the first tried backend outside the :obj:`set_backend` context manager. Note that this method is not thread-safe. .. warning:: We caution library authors against using this function in their code. We do *not* support this use-case. This function is meant to be used only by users themselves, or by a reference implementation, if one exists. Parameters ---------- backend The backend to register. coerce : bool Whether to coerce input types when trying this backend. only : bool If ``True``, no more backends will be tried if this fails. Implied by ``coerce=True``. try_last : bool If ``True``, the global backend is tried after registered backends. See Also -------- set_backend: A context manager that allows setting of backends. skip_backend: A context manager that allows skipping of backends. """ _uarray.set_global_backend(backend, coerce, only, try_last)
[docs]def register_backend(backend): """ This utility method sets registers backend for permanent use. It will be tried in the list of backends automatically, unless the ``only`` flag is set on a backend. Note that this method is not thread-safe. Parameters ---------- backend The backend to register. """ _uarray.register_backend(backend)
[docs]def clear_backends(domain, registered=True, globals=False): """ This utility method clears registered backends. .. warning:: We caution library authors against using this function in their code. We do *not* support this use-case. This function is meant to be used only by users themselves. .. warning:: Do NOT use this method inside a multimethod call, or the program is likely to crash. Parameters ---------- domain : Optional[str] The domain for which to de-register backends. ``None`` means de-register for all domains. registered : bool Whether or not to clear registered backends. See :obj:`register_backend`. globals : bool Whether or not to clear global backends. See :obj:`set_global_backend`. See Also -------- register_backend : Register a backend globally. set_global_backend : Set a global backend. """ _uarray.clear_backends(domain, registered, globals)
[docs]class Dispatchable: """ A utility class which marks an argument with a specific dispatch type. Attributes ---------- value The value of the Dispatchable. type The type of the Dispatchable. Examples -------- >>> x = Dispatchable(1, str) >>> x <Dispatchable: type=<class 'str'>, value=1> See Also -------- all_of_type Marks all unmarked parameters of a function. mark_as Allows one to create a utility function to mark as a given type. """
[docs] def __init__(self, value, dispatch_type, coercible=True): self.value = value self.type = dispatch_type self.coercible = coercible
def __getitem__(self, index): return (self.type, self.value)[index] def __str__(self): return "<{0}: type={1!r}, value={2!r}>".format( type(self).__name__, self.type, self.value ) __repr__ = __str__
[docs]def mark_as(dispatch_type): """ Creates a utility function to mark something as a specific type. Examples -------- >>> mark_int = mark_as(int) >>> mark_int(1) <Dispatchable: type=<class 'int'>, value=1> """ return functools.partial(Dispatchable, dispatch_type=dispatch_type)
[docs]def all_of_type(arg_type): """ Marks all unmarked arguments as a given type. Examples -------- >>> @all_of_type(str) ... def f(a, b): ... return a, Dispatchable(b, int) >>> f('a', 1) (<Dispatchable: type=<class 'str'>, value='a'>, <Dispatchable: type=<class 'int'>, value=1>) """ def outer(func): @functools.wraps(func) def inner(*args, **kwargs): extracted_args = func(*args, **kwargs) return tuple( Dispatchable(arg, arg_type) if not isinstance(arg, Dispatchable) else arg for arg in extracted_args ) return inner return outer
[docs]def wrap_single_convertor(convert_single): """ Wraps a ``__ua_convert__`` defined for a single element to all elements. If any of them return ``NotImplemented``, the operation is assumed to be undefined. Accepts a signature of (value, type, coerce). """ @functools.wraps(convert_single) def __ua_convert__(dispatchables, coerce): converted = [] for d in dispatchables: c = convert_single(d.value, d.type, coerce and d.coercible) if c is NotImplemented: return NotImplemented converted.append(c) return converted return __ua_convert__
def wrap_single_convertor_instance(convert_single): """ Wraps a ``__ua_convert__`` defined for a single element to all elements. If any of them return ``NotImplemented``, the operation is assumed to be undefined. Accepts a signature of (value, type, coerce). """ @functools.wraps(convert_single) def __ua_convert__(self, dispatchables, coerce): converted = [] for d in dispatchables: c = convert_single(self, d.value, d.type, coerce and d.coercible) if c is NotImplemented: return NotImplemented converted.append(c) return converted return __ua_convert__
[docs]def determine_backend(value, dispatch_type, *, domain, only=True, coerce=False): """Set the backend to the first active backend that supports ``value`` This is useful for functions that call multimethods without any dispatchable arguments. You can use :func:`determine_backend` to ensure the same backend is used everywhere in a block of multimethod calls. Parameters ---------- value The value being tested dispatch_type The dispatch type associated with ``value``, aka ":ref:`marking <MarkingGlossary>`". domain: string The domain to query for backends and set. coerce: bool Whether or not to allow coercion to the backend's types. Implies ``only``. only: bool Whether or not this should be the last backend to try. See Also -------- set_backend: For when you know which backend to set Notes ----- Support is determined by the ``__ua_convert__`` protocol. Backends not supporting the type must return ``NotImplemented`` from their ``__ua_convert__`` if they don't support input of that type. Examples -------- Suppose we have two backends ``BackendA`` and ``BackendB`` each supporting different types, ``TypeA`` and ``TypeB``. Neither supporting the other type: >>> with ua.set_backend(ex.BackendA): ... ex.call_multimethod(ex.TypeB(), ex.TypeB()) Traceback (most recent call last): ... uarray.BackendNotImplementedError: ... Now consider a multimethod that creates a new object of ``TypeA``, or ``TypeB`` depending on the active backend. >>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB): ... res = ex.creation_multimethod() ... ex.call_multimethod(res, ex.TypeA()) Traceback (most recent call last): ... uarray.BackendNotImplementedError: ... ``res`` is an object of ``TypeB`` because ``BackendB`` is set in the innermost with statement. So, ``call_multimethod`` fails since the types don't match. Instead, we need to first find a backend suitable for all of our objects. >>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB): ... x = ex.TypeA() ... with ua.determine_backend(x, "mark", domain="ua_examples"): ... res = ex.creation_multimethod() ... ex.call_multimethod(res, x) TypeA """ dispatchables = (Dispatchable(value, dispatch_type, coerce),) backend = _uarray.determine_backend(domain, dispatchables, coerce) return set_backend(backend, coerce=coerce, only=only)
[docs]def determine_backend_multi( dispatchables, *, domain, only=True, coerce=False, **kwargs ): """Set a backend supporting all ``dispatchables`` This is useful for functions that call multimethods without any dispatchable arguments. You can use :func:`determine_backend_multi` to ensure the same backend is used everywhere in a block of multimethod calls involving multiple arrays. Parameters ---------- dispatchables: Sequence[Union[uarray.Dispatchable, Any]] The dispatchables that must be supported domain: string The domain to query for backends and set. coerce: bool Whether or not to allow coercion to the backend's types. Implies ``only``. only: bool Whether or not this should be the last backend to try. dispatch_type: Optional[Any] The default dispatch type associated with ``dispatchables``, aka ":ref:`marking <MarkingGlossary>`". See Also -------- determine_backend: For a single dispatch value set_backend: For when you know which backend to set Notes ----- Support is determined by the ``__ua_convert__`` protocol. Backends not supporting the type must return ``NotImplemented`` from their ``__ua_convert__`` if they don't support input of that type. Examples -------- :func:`determine_backend` allows the backend to be set from a single object. :func:`determine_backend_multi` allows multiple objects to be checked simultaneously for support in the backend. Suppose we have a ``BackendAB`` which supports ``TypeA`` and ``TypeB`` in the same call, and a ``BackendBC`` that doesn't support ``TypeA``. >>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC): ... a, b = ex.TypeA(), ex.TypeB() ... with ua.determine_backend_multi( ... [ua.Dispatchable(a, "mark"), ua.Dispatchable(b, "mark")], ... domain="ua_examples" ... ): ... res = ex.creation_multimethod() ... ex.call_multimethod(res, a, b) TypeA This won't call ``BackendBC`` because it doesn't support ``TypeA``. We can also use leave out the ``ua.Dispatchable`` if we specify the default ``dispatch_type`` for the ``dispatchables`` argument. >>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC): ... a, b = ex.TypeA(), ex.TypeB() ... with ua.determine_backend_multi( ... [a, b], dispatch_type="mark", domain="ua_examples" ... ): ... res = ex.creation_multimethod() ... ex.call_multimethod(res, a, b) TypeA """ if "dispatch_type" in kwargs: disp_type = kwargs.pop("dispatch_type") dispatchables = tuple( d if isinstance(d, Dispatchable) else Dispatchable(d, disp_type) for d in dispatchables ) else: dispatchables = tuple(dispatchables) if not all(isinstance(d, Dispatchable) for d in dispatchables): raise TypeError("dispatchables must be instances of uarray.Dispatchable") if len(kwargs) != 0: raise TypeError("Received unexpected keyword arguments: {}".format(kwargs)) backend = _uarray.determine_backend(domain, dispatchables, coerce) return set_backend(backend, coerce=coerce, only=only)