Coverage for /usr/lib/python3/dist-packages/scipy/_lib/_uarray/_backend.py: 37%
163 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import typing
2import types
3import inspect
4import functools
5from . import _uarray
6import copyreg
7import pickle
8import contextlib
10from ._uarray import ( # type: ignore
11 BackendNotImplementedError,
12 _Function,
13 _SkipBackendContext,
14 _SetBackendContext,
15 _BackendState,
16)
18__all__ = [
19 "set_backend",
20 "set_global_backend",
21 "skip_backend",
22 "register_backend",
23 "determine_backend",
24 "determine_backend_multi",
25 "clear_backends",
26 "create_multimethod",
27 "generate_multimethod",
28 "_Function",
29 "BackendNotImplementedError",
30 "Dispatchable",
31 "wrap_single_convertor",
32 "wrap_single_convertor_instance",
33 "all_of_type",
34 "mark_as",
35 "set_state",
36 "get_state",
37 "reset_state",
38 "_BackendState",
39 "_SkipBackendContext",
40 "_SetBackendContext",
41]
43ArgumentExtractorType = typing.Callable[..., typing.Tuple["Dispatchable", ...]]
44ArgumentReplacerType = typing.Callable[
45 [typing.Tuple, typing.Dict, typing.Tuple], typing.Tuple[typing.Tuple, typing.Dict]
46]
48def unpickle_function(mod_name, qname, self_):
49 import importlib
51 try:
52 module = importlib.import_module(mod_name)
53 qname = qname.split(".")
54 func = module
55 for q in qname:
56 func = getattr(func, q)
58 if self_ is not None:
59 func = types.MethodType(func, self_)
61 return func
62 except (ImportError, AttributeError) as e:
63 from pickle import UnpicklingError
65 raise UnpicklingError from e
68def pickle_function(func):
69 mod_name = getattr(func, "__module__", None)
70 qname = getattr(func, "__qualname__", None)
71 self_ = getattr(func, "__self__", None)
73 try:
74 test = unpickle_function(mod_name, qname, self_)
75 except pickle.UnpicklingError:
76 test = None
78 if test is not func:
79 raise pickle.PicklingError(
80 f"Can't pickle {func}: it's not the same object as {test}"
81 )
83 return unpickle_function, (mod_name, qname, self_)
86def pickle_state(state):
87 return _uarray._BackendState._unpickle, state._pickle()
90def pickle_set_backend_context(ctx):
91 return _SetBackendContext, ctx._pickle()
94def pickle_skip_backend_context(ctx):
95 return _SkipBackendContext, ctx._pickle()
98copyreg.pickle(_Function, pickle_function)
99copyreg.pickle(_uarray._BackendState, pickle_state)
100copyreg.pickle(_SetBackendContext, pickle_set_backend_context)
101copyreg.pickle(_SkipBackendContext, pickle_skip_backend_context)
104def get_state():
105 """
106 Returns an opaque object containing the current state of all the backends.
108 Can be used for synchronization between threads/processes.
110 See Also
111 --------
112 set_state
113 Sets the state returned by this function.
114 """
115 return _uarray.get_state()
118@contextlib.contextmanager
119def reset_state():
120 """
121 Returns a context manager that resets all state once exited.
123 See Also
124 --------
125 set_state
126 Context manager that sets the backend state.
127 get_state
128 Gets a state to be set by this context manager.
129 """
130 with set_state(get_state()):
131 yield
134@contextlib.contextmanager
135def set_state(state):
136 """
137 A context manager that sets the state of the backends to one returned by :obj:`get_state`.
139 See Also
140 --------
141 get_state
142 Gets a state to be set by this context manager.
143 """
144 old_state = get_state()
145 _uarray.set_state(state)
146 try:
147 yield
148 finally:
149 _uarray.set_state(old_state, True)
152def create_multimethod(*args, **kwargs):
153 """
154 Creates a decorator for generating multimethods.
156 This function creates a decorator that can be used with an argument
157 extractor in order to generate a multimethod. Other than for the
158 argument extractor, all arguments are passed on to
159 :obj:`generate_multimethod`.
161 See Also
162 --------
163 generate_multimethod
164 Generates a multimethod.
165 """
167 def wrapper(a):
168 return generate_multimethod(a, *args, **kwargs)
170 return wrapper
173def generate_multimethod(
174 argument_extractor: ArgumentExtractorType,
175 argument_replacer: ArgumentReplacerType,
176 domain: str,
177 default: typing.Optional[typing.Callable] = None,
178):
179 """
180 Generates a multimethod.
182 Parameters
183 ----------
184 argument_extractor : ArgumentExtractorType
185 A callable which extracts the dispatchable arguments. Extracted arguments
186 should be marked by the :obj:`Dispatchable` class. It has the same signature
187 as the desired multimethod.
188 argument_replacer : ArgumentReplacerType
189 A callable with the signature (args, kwargs, dispatchables), which should also
190 return an (args, kwargs) pair with the dispatchables replaced inside the args/kwargs.
191 domain : str
192 A string value indicating the domain of this multimethod.
193 default: Optional[Callable], optional
194 The default implementation of this multimethod, where ``None`` (the default) specifies
195 there is no default implementation.
197 Examples
198 --------
199 In this example, ``a`` is to be dispatched over, so we return it, while marking it as an ``int``.
200 The trailing comma is needed because the args have to be returned as an iterable.
202 >>> def override_me(a, b):
203 ... return Dispatchable(a, int),
205 Next, we define the argument replacer that replaces the dispatchables inside args/kwargs with the
206 supplied ones.
208 >>> def override_replacer(args, kwargs, dispatchables):
209 ... return (dispatchables[0], args[1]), {}
211 Next, we define the multimethod.
213 >>> overridden_me = generate_multimethod(
214 ... override_me, override_replacer, "ua_examples"
215 ... )
217 Notice that there's no default implementation, unless you supply one.
219 >>> overridden_me(1, "a")
220 Traceback (most recent call last):
221 ...
222 uarray.BackendNotImplementedError: ...
224 >>> overridden_me2 = generate_multimethod(
225 ... override_me, override_replacer, "ua_examples", default=lambda x, y: (x, y)
226 ... )
227 >>> overridden_me2(1, "a")
228 (1, 'a')
230 See Also
231 --------
232 uarray
233 See the module documentation for how to override the method by creating backends.
234 """
235 kw_defaults, arg_defaults, opts = get_defaults(argument_extractor)
236 ua_func = _Function(
237 argument_extractor,
238 argument_replacer,
239 domain,
240 arg_defaults,
241 kw_defaults,
242 default,
243 )
245 return functools.update_wrapper(ua_func, argument_extractor)
248def set_backend(backend, coerce=False, only=False):
249 """
250 A context manager that sets the preferred backend.
252 Parameters
253 ----------
254 backend
255 The backend to set.
256 coerce
257 Whether or not to coerce to a specific backend's types. Implies ``only``.
258 only
259 Whether or not this should be the last backend to try.
261 See Also
262 --------
263 skip_backend: A context manager that allows skipping of backends.
264 set_global_backend: Set a single, global backend for a domain.
265 """
266 try:
267 return backend.__ua_cache__["set", coerce, only]
268 except AttributeError:
269 backend.__ua_cache__ = {}
270 except KeyError:
271 pass
273 ctx = _SetBackendContext(backend, coerce, only)
274 backend.__ua_cache__["set", coerce, only] = ctx
275 return ctx
278def skip_backend(backend):
279 """
280 A context manager that allows one to skip a given backend from processing
281 entirely. This allows one to use another backend's code in a library that
282 is also a consumer of the same backend.
284 Parameters
285 ----------
286 backend
287 The backend to skip.
289 See Also
290 --------
291 set_backend: A context manager that allows setting of backends.
292 set_global_backend: Set a single, global backend for a domain.
293 """
294 try:
295 return backend.__ua_cache__["skip"]
296 except AttributeError:
297 backend.__ua_cache__ = {}
298 except KeyError:
299 pass
301 ctx = _SkipBackendContext(backend)
302 backend.__ua_cache__["skip"] = ctx
303 return ctx
306def get_defaults(f):
307 sig = inspect.signature(f)
308 kw_defaults = {}
309 arg_defaults = []
310 opts = set()
311 for k, v in sig.parameters.items():
312 if v.default is not inspect.Parameter.empty:
313 kw_defaults[k] = v.default
314 if v.kind in (
315 inspect.Parameter.POSITIONAL_ONLY,
316 inspect.Parameter.POSITIONAL_OR_KEYWORD,
317 ):
318 arg_defaults.append(v.default)
319 opts.add(k)
321 return kw_defaults, tuple(arg_defaults), opts
324def set_global_backend(backend, coerce=False, only=False, *, try_last=False):
325 """
326 This utility method replaces the default backend for permanent use. It
327 will be tried in the list of backends automatically, unless the
328 ``only`` flag is set on a backend. This will be the first tried
329 backend outside the :obj:`set_backend` context manager.
331 Note that this method is not thread-safe.
333 .. warning::
334 We caution library authors against using this function in
335 their code. We do *not* support this use-case. This function
336 is meant to be used only by users themselves, or by a reference
337 implementation, if one exists.
339 Parameters
340 ----------
341 backend
342 The backend to register.
343 coerce : bool
344 Whether to coerce input types when trying this backend.
345 only : bool
346 If ``True``, no more backends will be tried if this fails.
347 Implied by ``coerce=True``.
348 try_last : bool
349 If ``True``, the global backend is tried after registered backends.
351 See Also
352 --------
353 set_backend: A context manager that allows setting of backends.
354 skip_backend: A context manager that allows skipping of backends.
355 """
356 _uarray.set_global_backend(backend, coerce, only, try_last)
359def register_backend(backend):
360 """
361 This utility method sets registers backend for permanent use. It
362 will be tried in the list of backends automatically, unless the
363 ``only`` flag is set on a backend.
365 Note that this method is not thread-safe.
367 Parameters
368 ----------
369 backend
370 The backend to register.
371 """
372 _uarray.register_backend(backend)
375def clear_backends(domain, registered=True, globals=False):
376 """
377 This utility method clears registered backends.
379 .. warning::
380 We caution library authors against using this function in
381 their code. We do *not* support this use-case. This function
382 is meant to be used only by users themselves.
384 .. warning::
385 Do NOT use this method inside a multimethod call, or the
386 program is likely to crash.
388 Parameters
389 ----------
390 domain : Optional[str]
391 The domain for which to de-register backends. ``None`` means
392 de-register for all domains.
393 registered : bool
394 Whether or not to clear registered backends. See :obj:`register_backend`.
395 globals : bool
396 Whether or not to clear global backends. See :obj:`set_global_backend`.
398 See Also
399 --------
400 register_backend : Register a backend globally.
401 set_global_backend : Set a global backend.
402 """
403 _uarray.clear_backends(domain, registered, globals)
406class Dispatchable:
407 """
408 A utility class which marks an argument with a specific dispatch type.
411 Attributes
412 ----------
413 value
414 The value of the Dispatchable.
416 type
417 The type of the Dispatchable.
419 Examples
420 --------
421 >>> x = Dispatchable(1, str)
422 >>> x
423 <Dispatchable: type=<class 'str'>, value=1>
425 See Also
426 --------
427 all_of_type
428 Marks all unmarked parameters of a function.
430 mark_as
431 Allows one to create a utility function to mark as a given type.
432 """
434 def __init__(self, value, dispatch_type, coercible=True):
435 self.value = value
436 self.type = dispatch_type
437 self.coercible = coercible
439 def __getitem__(self, index):
440 return (self.type, self.value)[index]
442 def __str__(self):
443 return "<{}: type={!r}, value={!r}>".format(
444 type(self).__name__, self.type, self.value
445 )
447 __repr__ = __str__
450def mark_as(dispatch_type):
451 """
452 Creates a utility function to mark something as a specific type.
454 Examples
455 --------
456 >>> mark_int = mark_as(int)
457 >>> mark_int(1)
458 <Dispatchable: type=<class 'int'>, value=1>
459 """
460 return functools.partial(Dispatchable, dispatch_type=dispatch_type)
463def all_of_type(arg_type):
464 """
465 Marks all unmarked arguments as a given type.
467 Examples
468 --------
469 >>> @all_of_type(str)
470 ... def f(a, b):
471 ... return a, Dispatchable(b, int)
472 >>> f('a', 1)
473 (<Dispatchable: type=<class 'str'>, value='a'>, <Dispatchable: type=<class 'int'>, value=1>)
474 """
476 def outer(func):
477 @functools.wraps(func)
478 def inner(*args, **kwargs):
479 extracted_args = func(*args, **kwargs)
480 return tuple(
481 Dispatchable(arg, arg_type)
482 if not isinstance(arg, Dispatchable)
483 else arg
484 for arg in extracted_args
485 )
487 return inner
489 return outer
492def wrap_single_convertor(convert_single):
493 """
494 Wraps a ``__ua_convert__`` defined for a single element to all elements.
495 If any of them return ``NotImplemented``, the operation is assumed to be
496 undefined.
498 Accepts a signature of (value, type, coerce).
499 """
501 @functools.wraps(convert_single)
502 def __ua_convert__(dispatchables, coerce):
503 converted = []
504 for d in dispatchables:
505 c = convert_single(d.value, d.type, coerce and d.coercible)
507 if c is NotImplemented:
508 return NotImplemented
510 converted.append(c)
512 return converted
514 return __ua_convert__
517def wrap_single_convertor_instance(convert_single):
518 """
519 Wraps a ``__ua_convert__`` defined for a single element to all elements.
520 If any of them return ``NotImplemented``, the operation is assumed to be
521 undefined.
523 Accepts a signature of (value, type, coerce).
524 """
526 @functools.wraps(convert_single)
527 def __ua_convert__(self, dispatchables, coerce):
528 converted = []
529 for d in dispatchables:
530 c = convert_single(self, d.value, d.type, coerce and d.coercible)
532 if c is NotImplemented:
533 return NotImplemented
535 converted.append(c)
537 return converted
539 return __ua_convert__
542def determine_backend(value, dispatch_type, *, domain, only=True, coerce=False):
543 """Set the backend to the first active backend that supports ``value``
545 This is useful for functions that call multimethods without any dispatchable
546 arguments. You can use :func:`determine_backend` to ensure the same backend
547 is used everywhere in a block of multimethod calls.
549 Parameters
550 ----------
551 value
552 The value being tested
553 dispatch_type
554 The dispatch type associated with ``value``, aka
555 ":ref:`marking <MarkingGlossary>`".
556 domain: string
557 The domain to query for backends and set.
558 coerce: bool
559 Whether or not to allow coercion to the backend's types. Implies ``only``.
560 only: bool
561 Whether or not this should be the last backend to try.
563 See Also
564 --------
565 set_backend: For when you know which backend to set
567 Notes
568 -----
570 Support is determined by the ``__ua_convert__`` protocol. Backends not
571 supporting the type must return ``NotImplemented`` from their
572 ``__ua_convert__`` if they don't support input of that type.
574 Examples
575 --------
577 Suppose we have two backends ``BackendA`` and ``BackendB`` each supporting
578 different types, ``TypeA`` and ``TypeB``. Neither supporting the other type:
580 >>> with ua.set_backend(ex.BackendA):
581 ... ex.call_multimethod(ex.TypeB(), ex.TypeB())
582 Traceback (most recent call last):
583 ...
584 uarray.BackendNotImplementedError: ...
586 Now consider a multimethod that creates a new object of ``TypeA``, or
587 ``TypeB`` depending on the active backend.
589 >>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
590 ... res = ex.creation_multimethod()
591 ... ex.call_multimethod(res, ex.TypeA())
592 Traceback (most recent call last):
593 ...
594 uarray.BackendNotImplementedError: ...
596 ``res`` is an object of ``TypeB`` because ``BackendB`` is set in the
597 innermost with statement. So, ``call_multimethod`` fails since the types
598 don't match.
600 Instead, we need to first find a backend suitable for all of our objects.
602 >>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
603 ... x = ex.TypeA()
604 ... with ua.determine_backend(x, "mark", domain="ua_examples"):
605 ... res = ex.creation_multimethod()
606 ... ex.call_multimethod(res, x)
607 TypeA
609 """
610 dispatchables = (Dispatchable(value, dispatch_type, coerce),)
611 backend = _uarray.determine_backend(domain, dispatchables, coerce)
613 return set_backend(backend, coerce=coerce, only=only)
616def determine_backend_multi(
617 dispatchables, *, domain, only=True, coerce=False, **kwargs
618):
619 """Set a backend supporting all ``dispatchables``
621 This is useful for functions that call multimethods without any dispatchable
622 arguments. You can use :func:`determine_backend_multi` to ensure the same
623 backend is used everywhere in a block of multimethod calls involving
624 multiple arrays.
626 Parameters
627 ----------
628 dispatchables: Sequence[Union[uarray.Dispatchable, Any]]
629 The dispatchables that must be supported
630 domain: string
631 The domain to query for backends and set.
632 coerce: bool
633 Whether or not to allow coercion to the backend's types. Implies ``only``.
634 only: bool
635 Whether or not this should be the last backend to try.
636 dispatch_type: Optional[Any]
637 The default dispatch type associated with ``dispatchables``, aka
638 ":ref:`marking <MarkingGlossary>`".
640 See Also
641 --------
642 determine_backend: For a single dispatch value
643 set_backend: For when you know which backend to set
645 Notes
646 -----
648 Support is determined by the ``__ua_convert__`` protocol. Backends not
649 supporting the type must return ``NotImplemented`` from their
650 ``__ua_convert__`` if they don't support input of that type.
652 Examples
653 --------
655 :func:`determine_backend` allows the backend to be set from a single
656 object. :func:`determine_backend_multi` allows multiple objects to be
657 checked simultaneously for support in the backend. Suppose we have a
658 ``BackendAB`` which supports ``TypeA`` and ``TypeB`` in the same call,
659 and a ``BackendBC`` that doesn't support ``TypeA``.
661 >>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
662 ... a, b = ex.TypeA(), ex.TypeB()
663 ... with ua.determine_backend_multi(
664 ... [ua.Dispatchable(a, "mark"), ua.Dispatchable(b, "mark")],
665 ... domain="ua_examples"
666 ... ):
667 ... res = ex.creation_multimethod()
668 ... ex.call_multimethod(res, a, b)
669 TypeA
671 This won't call ``BackendBC`` because it doesn't support ``TypeA``.
673 We can also use leave out the ``ua.Dispatchable`` if we specify the
674 default ``dispatch_type`` for the ``dispatchables`` argument.
676 >>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
677 ... a, b = ex.TypeA(), ex.TypeB()
678 ... with ua.determine_backend_multi(
679 ... [a, b], dispatch_type="mark", domain="ua_examples"
680 ... ):
681 ... res = ex.creation_multimethod()
682 ... ex.call_multimethod(res, a, b)
683 TypeA
685 """
686 if "dispatch_type" in kwargs:
687 disp_type = kwargs.pop("dispatch_type")
688 dispatchables = tuple(
689 d if isinstance(d, Dispatchable) else Dispatchable(d, disp_type)
690 for d in dispatchables
691 )
692 else:
693 dispatchables = tuple(dispatchables)
694 if not all(isinstance(d, Dispatchable) for d in dispatchables):
695 raise TypeError("dispatchables must be instances of uarray.Dispatchable")
697 if len(kwargs) != 0:
698 raise TypeError(f"Received unexpected keyword arguments: {kwargs}")
700 backend = _uarray.determine_backend(domain, dispatchables, coerce)
702 return set_backend(backend, coerce=coerce, only=only)