Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/scipy/_lib/array_api_compat/_internal.py: 93%
14 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-03 06:39 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-03 06:39 +0000
1"""
2Internal helpers
3"""
5from functools import wraps
6from inspect import signature
8def get_xp(xp):
9 """
10 Decorator to automatically replace xp with the corresponding array module.
12 Use like
14 import numpy as np
16 @get_xp(np)
17 def func(x, /, xp, kwarg=None):
18 return xp.func(x, kwarg=kwarg)
20 Note that xp must be a keyword argument and come after all non-keyword
21 arguments.
23 """
25 def inner(f):
26 @wraps(f)
27 def wrapped_f(*args, **kwargs):
28 return f(*args, xp=xp, **kwargs)
30 sig = signature(f)
31 new_sig = sig.replace(
32 parameters=[sig.parameters[i] for i in sig.parameters if i != "xp"]
33 )
35 if wrapped_f.__doc__ is None:
36 wrapped_f.__doc__ = f"""\
37Array API compatibility wrapper for {f.__name__}.
39See the corresponding documentation in NumPy/CuPy and/or the array API
40specification for more details.
42"""
43 wrapped_f.__signature__ = new_sig
44 return wrapped_f
46 return inner