Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1"""Implementation of __array_function__ overrides from NEP-18.""" 

2import collections 

3import functools 

4import os 

5import textwrap 

6 

7from numpy.core._multiarray_umath import ( 

8 add_docstring, implement_array_function, _get_implementing_args) 

9from numpy.compat._inspect import getargspec 

10 

11 

12ARRAY_FUNCTION_ENABLED = bool( 

13 int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1))) 

14 

15 

16add_docstring( 

17 implement_array_function, 

18 """ 

19 Implement a function with checks for __array_function__ overrides. 

20 

21 All arguments are required, and can only be passed by position. 

22 

23 Arguments 

24 --------- 

25 implementation : function 

26 Function that implements the operation on NumPy array without 

27 overrides when called like ``implementation(*args, **kwargs)``. 

28 public_api : function 

29 Function exposed by NumPy's public API originally called like 

30 ``public_api(*args, **kwargs)`` on which arguments are now being 

31 checked. 

32 relevant_args : iterable 

33 Iterable of arguments to check for __array_function__ methods. 

34 args : tuple 

35 Arbitrary positional arguments originally passed into ``public_api``. 

36 kwargs : dict 

37 Arbitrary keyword arguments originally passed into ``public_api``. 

38 

39 Returns 

40 ------- 

41 Result from calling ``implementation()`` or an ``__array_function__`` 

42 method, as appropriate. 

43 

44 Raises 

45 ------ 

46 TypeError : if no implementation is found. 

47 """) 

48 

49 

50# exposed for testing purposes; used internally by implement_array_function 

51add_docstring( 

52 _get_implementing_args, 

53 """ 

54 Collect arguments on which to call __array_function__. 

55 

56 Parameters 

57 ---------- 

58 relevant_args : iterable of array-like 

59 Iterable of possibly array-like arguments to check for 

60 __array_function__ methods. 

61 

62 Returns 

63 ------- 

64 Sequence of arguments with __array_function__ methods, in the order in 

65 which they should be called. 

66 """) 

67 

68 

69ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults') 

70 

71 

72def verify_matching_signatures(implementation, dispatcher): 

73 """Verify that a dispatcher function has the right signature.""" 

74 implementation_spec = ArgSpec(*getargspec(implementation)) 

75 dispatcher_spec = ArgSpec(*getargspec(dispatcher)) 

76 

77 if (implementation_spec.args != dispatcher_spec.args or 

78 implementation_spec.varargs != dispatcher_spec.varargs or 

79 implementation_spec.keywords != dispatcher_spec.keywords or 

80 (bool(implementation_spec.defaults) != 

81 bool(dispatcher_spec.defaults)) or 

82 (implementation_spec.defaults is not None and 

83 len(implementation_spec.defaults) != 

84 len(dispatcher_spec.defaults))): 

85 raise RuntimeError('implementation and dispatcher for %s have ' 

86 'different function signatures' % implementation) 

87 

88 if implementation_spec.defaults is not None: 

89 if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults): 

90 raise RuntimeError('dispatcher functions can only use None for ' 

91 'default argument values') 

92 

93 

94def set_module(module): 

95 """Decorator for overriding __module__ on a function or class. 

96 

97 Example usage:: 

98 

99 @set_module('numpy') 

100 def example(): 

101 pass 

102 

103 assert example.__module__ == 'numpy' 

104 """ 

105 def decorator(func): 

106 if module is not None: 

107 func.__module__ = module 

108 return func 

109 return decorator 

110 

111 

112 

113# Call textwrap.dedent here instead of in the function so as to avoid 

114# calling dedent multiple times on the same text 

115_wrapped_func_source = textwrap.dedent(""" 

116 @functools.wraps(implementation) 

117 def {name}(*args, **kwargs): 

118 relevant_args = dispatcher(*args, **kwargs) 

119 return implement_array_function( 

120 implementation, {name}, relevant_args, args, kwargs) 

121 """) 

122 

123 

124def array_function_dispatch(dispatcher, module=None, verify=True, 

125 docs_from_dispatcher=False): 

126 """Decorator for adding dispatch with the __array_function__ protocol. 

127 

128 See NEP-18 for example usage. 

129 

130 Parameters 

131 ---------- 

132 dispatcher : callable 

133 Function that when called like ``dispatcher(*args, **kwargs)`` with 

134 arguments from the NumPy function call returns an iterable of 

135 array-like arguments to check for ``__array_function__``. 

136 module : str, optional 

137 __module__ attribute to set on new function, e.g., ``module='numpy'``. 

138 By default, module is copied from the decorated function. 

139 verify : bool, optional 

140 If True, verify the that the signature of the dispatcher and decorated 

141 function signatures match exactly: all required and optional arguments 

142 should appear in order with the same names, but the default values for 

143 all optional arguments should be ``None``. Only disable verification 

144 if the dispatcher's signature needs to deviate for some particular 

145 reason, e.g., because the function has a signature like 

146 ``func(*args, **kwargs)``. 

147 docs_from_dispatcher : bool, optional 

148 If True, copy docs from the dispatcher function onto the dispatched 

149 function, rather than from the implementation. This is useful for 

150 functions defined in C, which otherwise don't have docstrings. 

151 

152 Returns 

153 ------- 

154 Function suitable for decorating the implementation of a NumPy function. 

155 """ 

156 

157 if not ARRAY_FUNCTION_ENABLED: 

158 def decorator(implementation): 

159 if docs_from_dispatcher: 

160 add_docstring(implementation, dispatcher.__doc__) 

161 if module is not None: 

162 implementation.__module__ = module 

163 return implementation 

164 return decorator 

165 

166 def decorator(implementation): 

167 if verify: 

168 verify_matching_signatures(implementation, dispatcher) 

169 

170 if docs_from_dispatcher: 

171 add_docstring(implementation, dispatcher.__doc__) 

172 

173 # Equivalently, we could define this function directly instead of using 

174 # exec. This version has the advantage of giving the helper function a 

175 # more interpettable name. Otherwise, the original function does not 

176 # show up at all in many cases, e.g., if it's written in C or if the 

177 # dispatcher gets an invalid keyword argument. 

178 source = _wrapped_func_source.format(name=implementation.__name__) 

179 

180 source_object = compile( 

181 source, filename='<__array_function__ internals>', mode='exec') 

182 scope = { 

183 'implementation': implementation, 

184 'dispatcher': dispatcher, 

185 'functools': functools, 

186 'implement_array_function': implement_array_function, 

187 } 

188 exec(source_object, scope) 

189 

190 public_api = scope[implementation.__name__] 

191 

192 if module is not None: 

193 public_api.__module__ = module 

194 

195 public_api._implementation = implementation 

196 

197 return public_api 

198 

199 return decorator 

200 

201 

202def array_function_from_dispatcher( 

203 implementation, module=None, verify=True, docs_from_dispatcher=True): 

204 """Like array_function_dispatcher, but with function arguments flipped.""" 

205 

206 def decorator(dispatcher): 

207 return array_function_dispatch( 

208 dispatcher, module, verify=verify, 

209 docs_from_dispatcher=docs_from_dispatcher)(implementation) 

210 return decorator