_internal.py

"""
Internal helpers
"""

fromcollections.abcimportCallable
fromfunctoolsimportwraps
frominspectimportsignature
fromtypesimportModuleType
fromtypingimportTypeVar

_T=TypeVar("_T")


defget_xp(xp: ModuleType) ->Callable[[Callable[..., _T]], Callable[..., _T]]:
"""
 Decorator to automatically replace xp with the corresponding array module.

 Use like

 import numpy as np

 @get_xp(np)
 def func(x, /, xp, kwarg=None):
 return xp.func(x, kwarg=kwarg)

 Note that xp must be a keyword argument and come after all non-keyword
 arguments.

 """

definner(f: Callable[..., _T], /) ->Callable[..., _T]:
@wraps(f)
defwrapped_f(*args: object, **kwargs: object) ->object:
returnf(*args, xp=xp, **kwargs)

sig=signature(f)
new_sig=sig.replace(
parameters=[parfori, parinsig.parameters.items() ifi!="xp"]
 )

ifwrapped_f.__doc__isNone:
wrapped_f.__doc__=f"""\
Array API compatibility wrapper for {f.__name__}.

See the corresponding documentation in NumPy/CuPy and/or the array API
specification for more details.

"""
wrapped_f.__signature__=new_sig# pyright: ignore[reportAttributeAccessIssue]
returnwrapped_f# pyright: ignore[reportReturnType]

returninner


__all__= ["get_xp"]


def__dir__() ->list[str]:
return__all__