overrides.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
  1. """Implementation of __array_function__ overrides from NEP-18."""
  2. import collections
  3. import functools
  4. import os
  5. import textwrap
  6. from numpy.core._multiarray_umath import (
  7. add_docstring, implement_array_function, _get_implementing_args)
  8. from numpy.compat._inspect import getargspec
  9. ARRAY_FUNCTION_ENABLED = bool(
  10. int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1)))
  11. add_docstring(
  12. implement_array_function,
  13. """
  14. Implement a function with checks for __array_function__ overrides.
  15. All arguments are required, and can only be passed by position.
  16. Arguments
  17. ---------
  18. implementation : function
  19. Function that implements the operation on NumPy array without
  20. overrides when called like ``implementation(*args, **kwargs)``.
  21. public_api : function
  22. Function exposed by NumPy's public API originally called like
  23. ``public_api(*args, **kwargs)`` on which arguments are now being
  24. checked.
  25. relevant_args : iterable
  26. Iterable of arguments to check for __array_function__ methods.
  27. args : tuple
  28. Arbitrary positional arguments originally passed into ``public_api``.
  29. kwargs : dict
  30. Arbitrary keyword arguments originally passed into ``public_api``.
  31. Returns
  32. -------
  33. Result from calling ``implementation()`` or an ``__array_function__``
  34. method, as appropriate.
  35. Raises
  36. ------
  37. TypeError : if no implementation is found.
  38. """)
  39. # exposed for testing purposes; used internally by implement_array_function
  40. add_docstring(
  41. _get_implementing_args,
  42. """
  43. Collect arguments on which to call __array_function__.
  44. Parameters
  45. ----------
  46. relevant_args : iterable of array-like
  47. Iterable of possibly array-like arguments to check for
  48. __array_function__ methods.
  49. Returns
  50. -------
  51. Sequence of arguments with __array_function__ methods, in the order in
  52. which they should be called.
  53. """)
  54. ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')
  55. def verify_matching_signatures(implementation, dispatcher):
  56. """Verify that a dispatcher function has the right signature."""
  57. implementation_spec = ArgSpec(*getargspec(implementation))
  58. dispatcher_spec = ArgSpec(*getargspec(dispatcher))
  59. if (implementation_spec.args != dispatcher_spec.args or
  60. implementation_spec.varargs != dispatcher_spec.varargs or
  61. implementation_spec.keywords != dispatcher_spec.keywords or
  62. (bool(implementation_spec.defaults) !=
  63. bool(dispatcher_spec.defaults)) or
  64. (implementation_spec.defaults is not None and
  65. len(implementation_spec.defaults) !=
  66. len(dispatcher_spec.defaults))):
  67. raise RuntimeError('implementation and dispatcher for %s have '
  68. 'different function signatures' % implementation)
  69. if implementation_spec.defaults is not None:
  70. if dispatcher_spec.defaults != (None,) * len(dispatcher_spec.defaults):
  71. raise RuntimeError('dispatcher functions can only use None for '
  72. 'default argument values')
  73. def set_module(module):
  74. """Decorator for overriding __module__ on a function or class.
  75. Example usage::
  76. @set_module('numpy')
  77. def example():
  78. pass
  79. assert example.__module__ == 'numpy'
  80. """
  81. def decorator(func):
  82. if module is not None:
  83. func.__module__ = module
  84. return func
  85. return decorator
  86. # Call textwrap.dedent here instead of in the function so as to avoid
  87. # calling dedent multiple times on the same text
  88. _wrapped_func_source = textwrap.dedent("""
  89. @functools.wraps(implementation)
  90. def {name}(*args, **kwargs):
  91. relevant_args = dispatcher(*args, **kwargs)
  92. return implement_array_function(
  93. implementation, {name}, relevant_args, args, kwargs)
  94. """)
  95. def array_function_dispatch(dispatcher, module=None, verify=True,
  96. docs_from_dispatcher=False):
  97. """Decorator for adding dispatch with the __array_function__ protocol.
  98. See NEP-18 for example usage.
  99. Parameters
  100. ----------
  101. dispatcher : callable
  102. Function that when called like ``dispatcher(*args, **kwargs)`` with
  103. arguments from the NumPy function call returns an iterable of
  104. array-like arguments to check for ``__array_function__``.
  105. module : str, optional
  106. __module__ attribute to set on new function, e.g., ``module='numpy'``.
  107. By default, module is copied from the decorated function.
  108. verify : bool, optional
  109. If True, verify the that the signature of the dispatcher and decorated
  110. function signatures match exactly: all required and optional arguments
  111. should appear in order with the same names, but the default values for
  112. all optional arguments should be ``None``. Only disable verification
  113. if the dispatcher's signature needs to deviate for some particular
  114. reason, e.g., because the function has a signature like
  115. ``func(*args, **kwargs)``.
  116. docs_from_dispatcher : bool, optional
  117. If True, copy docs from the dispatcher function onto the dispatched
  118. function, rather than from the implementation. This is useful for
  119. functions defined in C, which otherwise don't have docstrings.
  120. Returns
  121. -------
  122. Function suitable for decorating the implementation of a NumPy function.
  123. """
  124. if not ARRAY_FUNCTION_ENABLED:
  125. def decorator(implementation):
  126. if docs_from_dispatcher:
  127. add_docstring(implementation, dispatcher.__doc__)
  128. if module is not None:
  129. implementation.__module__ = module
  130. return implementation
  131. return decorator
  132. def decorator(implementation):
  133. if verify:
  134. verify_matching_signatures(implementation, dispatcher)
  135. if docs_from_dispatcher:
  136. add_docstring(implementation, dispatcher.__doc__)
  137. # Equivalently, we could define this function directly instead of using
  138. # exec. This version has the advantage of giving the helper function a
  139. # more interpettable name. Otherwise, the original function does not
  140. # show up at all in many cases, e.g., if it's written in C or if the
  141. # dispatcher gets an invalid keyword argument.
  142. source = _wrapped_func_source.format(name=implementation.__name__)
  143. source_object = compile(
  144. source, filename='<__array_function__ internals>', mode='exec')
  145. scope = {
  146. 'implementation': implementation,
  147. 'dispatcher': dispatcher,
  148. 'functools': functools,
  149. 'implement_array_function': implement_array_function,
  150. }
  151. exec(source_object, scope)
  152. public_api = scope[implementation.__name__]
  153. if module is not None:
  154. public_api.__module__ = module
  155. public_api._implementation = implementation
  156. return public_api
  157. return decorator
  158. def array_function_from_dispatcher(
  159. implementation, module=None, verify=True, docs_from_dispatcher=True):
  160. """Like array_function_dispatcher, but with function arguments flipped."""
  161. def decorator(dispatcher):
  162. return array_function_dispatch(
  163. dispatcher, module, verify=verify,
  164. docs_from_dispatcher=docs_from_dispatcher)(implementation)
  165. return decorator