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

1""" 

2Internal helpers 

3""" 

4 

5from functools import wraps 

6from inspect import signature 

7 

8def get_xp(xp): 

9 """ 

10 Decorator to automatically replace xp with the corresponding array module. 

11 

12 Use like 

13 

14 import numpy as np 

15 

16 @get_xp(np) 

17 def func(x, /, xp, kwarg=None): 

18 return xp.func(x, kwarg=kwarg) 

19 

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

21 arguments. 

22 

23 """ 

24 

25 def inner(f): 

26 @wraps(f) 

27 def wrapped_f(*args, **kwargs): 

28 return f(*args, xp=xp, **kwargs) 

29 

30 sig = signature(f) 

31 new_sig = sig.replace( 

32 parameters=[sig.parameters[i] for i in sig.parameters if i != "xp"] 

33 ) 

34 

35 if wrapped_f.__doc__ is None: 

36 wrapped_f.__doc__ = f"""\ 

37Array API compatibility wrapper for {f.__name__}. 

38 

39See the corresponding documentation in NumPy/CuPy and/or the array API 

40specification for more details. 

41 

42""" 

43 wrapped_f.__signature__ = new_sig 

44 return wrapped_f 

45 

46 return inner