test_overrides.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427
  1. import inspect
  2. import sys
  3. from unittest import mock
  4. import numpy as np
  5. from numpy.testing import (
  6. assert_, assert_equal, assert_raises, assert_raises_regex)
  7. from numpy.core.overrides import (
  8. _get_implementing_args, array_function_dispatch,
  9. verify_matching_signatures, ARRAY_FUNCTION_ENABLED)
  10. from numpy.compat import pickle
  11. import pytest
  12. requires_array_function = pytest.mark.skipif(
  13. not ARRAY_FUNCTION_ENABLED,
  14. reason="__array_function__ dispatch not enabled.")
  15. def _return_not_implemented(self, *args, **kwargs):
  16. return NotImplemented
  17. # need to define this at the top level to test pickling
  18. @array_function_dispatch(lambda array: (array,))
  19. def dispatched_one_arg(array):
  20. """Docstring."""
  21. return 'original'
  22. @array_function_dispatch(lambda array1, array2: (array1, array2))
  23. def dispatched_two_arg(array1, array2):
  24. """Docstring."""
  25. return 'original'
  26. class TestGetImplementingArgs:
  27. def test_ndarray(self):
  28. array = np.array(1)
  29. args = _get_implementing_args([array])
  30. assert_equal(list(args), [array])
  31. args = _get_implementing_args([array, array])
  32. assert_equal(list(args), [array])
  33. args = _get_implementing_args([array, 1])
  34. assert_equal(list(args), [array])
  35. args = _get_implementing_args([1, array])
  36. assert_equal(list(args), [array])
  37. def test_ndarray_subclasses(self):
  38. class OverrideSub(np.ndarray):
  39. __array_function__ = _return_not_implemented
  40. class NoOverrideSub(np.ndarray):
  41. pass
  42. array = np.array(1).view(np.ndarray)
  43. override_sub = np.array(1).view(OverrideSub)
  44. no_override_sub = np.array(1).view(NoOverrideSub)
  45. args = _get_implementing_args([array, override_sub])
  46. assert_equal(list(args), [override_sub, array])
  47. args = _get_implementing_args([array, no_override_sub])
  48. assert_equal(list(args), [no_override_sub, array])
  49. args = _get_implementing_args(
  50. [override_sub, no_override_sub])
  51. assert_equal(list(args), [override_sub, no_override_sub])
  52. def test_ndarray_and_duck_array(self):
  53. class Other:
  54. __array_function__ = _return_not_implemented
  55. array = np.array(1)
  56. other = Other()
  57. args = _get_implementing_args([other, array])
  58. assert_equal(list(args), [other, array])
  59. args = _get_implementing_args([array, other])
  60. assert_equal(list(args), [array, other])
  61. def test_ndarray_subclass_and_duck_array(self):
  62. class OverrideSub(np.ndarray):
  63. __array_function__ = _return_not_implemented
  64. class Other:
  65. __array_function__ = _return_not_implemented
  66. array = np.array(1)
  67. subarray = np.array(1).view(OverrideSub)
  68. other = Other()
  69. assert_equal(_get_implementing_args([array, subarray, other]),
  70. [subarray, array, other])
  71. assert_equal(_get_implementing_args([array, other, subarray]),
  72. [subarray, array, other])
  73. def test_many_duck_arrays(self):
  74. class A:
  75. __array_function__ = _return_not_implemented
  76. class B(A):
  77. __array_function__ = _return_not_implemented
  78. class C(A):
  79. __array_function__ = _return_not_implemented
  80. class D:
  81. __array_function__ = _return_not_implemented
  82. a = A()
  83. b = B()
  84. c = C()
  85. d = D()
  86. assert_equal(_get_implementing_args([1]), [])
  87. assert_equal(_get_implementing_args([a]), [a])
  88. assert_equal(_get_implementing_args([a, 1]), [a])
  89. assert_equal(_get_implementing_args([a, a, a]), [a])
  90. assert_equal(_get_implementing_args([a, d, a]), [a, d])
  91. assert_equal(_get_implementing_args([a, b]), [b, a])
  92. assert_equal(_get_implementing_args([b, a]), [b, a])
  93. assert_equal(_get_implementing_args([a, b, c]), [b, c, a])
  94. assert_equal(_get_implementing_args([a, c, b]), [c, b, a])
  95. def test_too_many_duck_arrays(self):
  96. namespace = dict(__array_function__=_return_not_implemented)
  97. types = [type('A' + str(i), (object,), namespace) for i in range(33)]
  98. relevant_args = [t() for t in types]
  99. actual = _get_implementing_args(relevant_args[:32])
  100. assert_equal(actual, relevant_args[:32])
  101. with assert_raises_regex(TypeError, 'distinct argument types'):
  102. _get_implementing_args(relevant_args)
  103. class TestNDArrayArrayFunction:
  104. @requires_array_function
  105. def test_method(self):
  106. class Other:
  107. __array_function__ = _return_not_implemented
  108. class NoOverrideSub(np.ndarray):
  109. pass
  110. class OverrideSub(np.ndarray):
  111. __array_function__ = _return_not_implemented
  112. array = np.array([1])
  113. other = Other()
  114. no_override_sub = array.view(NoOverrideSub)
  115. override_sub = array.view(OverrideSub)
  116. result = array.__array_function__(func=dispatched_two_arg,
  117. types=(np.ndarray,),
  118. args=(array, 1.), kwargs={})
  119. assert_equal(result, 'original')
  120. result = array.__array_function__(func=dispatched_two_arg,
  121. types=(np.ndarray, Other),
  122. args=(array, other), kwargs={})
  123. assert_(result is NotImplemented)
  124. result = array.__array_function__(func=dispatched_two_arg,
  125. types=(np.ndarray, NoOverrideSub),
  126. args=(array, no_override_sub),
  127. kwargs={})
  128. assert_equal(result, 'original')
  129. result = array.__array_function__(func=dispatched_two_arg,
  130. types=(np.ndarray, OverrideSub),
  131. args=(array, override_sub),
  132. kwargs={})
  133. assert_equal(result, 'original')
  134. with assert_raises_regex(TypeError, 'no implementation found'):
  135. np.concatenate((array, other))
  136. expected = np.concatenate((array, array))
  137. result = np.concatenate((array, no_override_sub))
  138. assert_equal(result, expected.view(NoOverrideSub))
  139. result = np.concatenate((array, override_sub))
  140. assert_equal(result, expected.view(OverrideSub))
  141. def test_no_wrapper(self):
  142. # This shouldn't happen unless a user intentionally calls
  143. # __array_function__ with invalid arguments, but check that we raise
  144. # an appropriate error all the same.
  145. array = np.array(1)
  146. func = lambda x: x
  147. with assert_raises_regex(AttributeError, '_implementation'):
  148. array.__array_function__(func=func, types=(np.ndarray,),
  149. args=(array,), kwargs={})
  150. @requires_array_function
  151. class TestArrayFunctionDispatch:
  152. def test_pickle(self):
  153. for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
  154. roundtripped = pickle.loads(
  155. pickle.dumps(dispatched_one_arg, protocol=proto))
  156. assert_(roundtripped is dispatched_one_arg)
  157. def test_name_and_docstring(self):
  158. assert_equal(dispatched_one_arg.__name__, 'dispatched_one_arg')
  159. if sys.flags.optimize < 2:
  160. assert_equal(dispatched_one_arg.__doc__, 'Docstring.')
  161. def test_interface(self):
  162. class MyArray:
  163. def __array_function__(self, func, types, args, kwargs):
  164. return (self, func, types, args, kwargs)
  165. original = MyArray()
  166. (obj, func, types, args, kwargs) = dispatched_one_arg(original)
  167. assert_(obj is original)
  168. assert_(func is dispatched_one_arg)
  169. assert_equal(set(types), {MyArray})
  170. # assert_equal uses the overloaded np.iscomplexobj() internally
  171. assert_(args == (original,))
  172. assert_equal(kwargs, {})
  173. def test_not_implemented(self):
  174. class MyArray:
  175. def __array_function__(self, func, types, args, kwargs):
  176. return NotImplemented
  177. array = MyArray()
  178. with assert_raises_regex(TypeError, 'no implementation found'):
  179. dispatched_one_arg(array)
  180. @requires_array_function
  181. class TestVerifyMatchingSignatures:
  182. def test_verify_matching_signatures(self):
  183. verify_matching_signatures(lambda x: 0, lambda x: 0)
  184. verify_matching_signatures(lambda x=None: 0, lambda x=None: 0)
  185. verify_matching_signatures(lambda x=1: 0, lambda x=None: 0)
  186. with assert_raises(RuntimeError):
  187. verify_matching_signatures(lambda a: 0, lambda b: 0)
  188. with assert_raises(RuntimeError):
  189. verify_matching_signatures(lambda x: 0, lambda x=None: 0)
  190. with assert_raises(RuntimeError):
  191. verify_matching_signatures(lambda x=None: 0, lambda y=None: 0)
  192. with assert_raises(RuntimeError):
  193. verify_matching_signatures(lambda x=1: 0, lambda y=1: 0)
  194. def test_array_function_dispatch(self):
  195. with assert_raises(RuntimeError):
  196. @array_function_dispatch(lambda x: (x,))
  197. def f(y):
  198. pass
  199. # should not raise
  200. @array_function_dispatch(lambda x: (x,), verify=False)
  201. def f(y):
  202. pass
  203. def _new_duck_type_and_implements():
  204. """Create a duck array type and implements functions."""
  205. HANDLED_FUNCTIONS = {}
  206. class MyArray:
  207. def __array_function__(self, func, types, args, kwargs):
  208. if func not in HANDLED_FUNCTIONS:
  209. return NotImplemented
  210. if not all(issubclass(t, MyArray) for t in types):
  211. return NotImplemented
  212. return HANDLED_FUNCTIONS[func](*args, **kwargs)
  213. def implements(numpy_function):
  214. """Register an __array_function__ implementations."""
  215. def decorator(func):
  216. HANDLED_FUNCTIONS[numpy_function] = func
  217. return func
  218. return decorator
  219. return (MyArray, implements)
  220. @requires_array_function
  221. class TestArrayFunctionImplementation:
  222. def test_one_arg(self):
  223. MyArray, implements = _new_duck_type_and_implements()
  224. @implements(dispatched_one_arg)
  225. def _(array):
  226. return 'myarray'
  227. assert_equal(dispatched_one_arg(1), 'original')
  228. assert_equal(dispatched_one_arg(MyArray()), 'myarray')
  229. def test_optional_args(self):
  230. MyArray, implements = _new_duck_type_and_implements()
  231. @array_function_dispatch(lambda array, option=None: (array,))
  232. def func_with_option(array, option='default'):
  233. return option
  234. @implements(func_with_option)
  235. def my_array_func_with_option(array, new_option='myarray'):
  236. return new_option
  237. # we don't need to implement every option on __array_function__
  238. # implementations
  239. assert_equal(func_with_option(1), 'default')
  240. assert_equal(func_with_option(1, option='extra'), 'extra')
  241. assert_equal(func_with_option(MyArray()), 'myarray')
  242. with assert_raises(TypeError):
  243. func_with_option(MyArray(), option='extra')
  244. # but new options on implementations can't be used
  245. result = my_array_func_with_option(MyArray(), new_option='yes')
  246. assert_equal(result, 'yes')
  247. with assert_raises(TypeError):
  248. func_with_option(MyArray(), new_option='no')
  249. def test_not_implemented(self):
  250. MyArray, implements = _new_duck_type_and_implements()
  251. @array_function_dispatch(lambda array: (array,), module='my')
  252. def func(array):
  253. return array
  254. array = np.array(1)
  255. assert_(func(array) is array)
  256. assert_equal(func.__module__, 'my')
  257. with assert_raises_regex(
  258. TypeError, "no implementation found for 'my.func'"):
  259. func(MyArray())
  260. class TestNDArrayMethods:
  261. def test_repr(self):
  262. # gh-12162: should still be defined even if __array_function__ doesn't
  263. # implement np.array_repr()
  264. class MyArray(np.ndarray):
  265. def __array_function__(*args, **kwargs):
  266. return NotImplemented
  267. array = np.array(1).view(MyArray)
  268. assert_equal(repr(array), 'MyArray(1)')
  269. assert_equal(str(array), '1')
  270. class TestNumPyFunctions:
  271. def test_set_module(self):
  272. assert_equal(np.sum.__module__, 'numpy')
  273. assert_equal(np.char.equal.__module__, 'numpy.char')
  274. assert_equal(np.fft.fft.__module__, 'numpy.fft')
  275. assert_equal(np.linalg.solve.__module__, 'numpy.linalg')
  276. def test_inspect_sum(self):
  277. signature = inspect.signature(np.sum)
  278. assert_('axis' in signature.parameters)
  279. @requires_array_function
  280. def test_override_sum(self):
  281. MyArray, implements = _new_duck_type_and_implements()
  282. @implements(np.sum)
  283. def _(array):
  284. return 'yes'
  285. assert_equal(np.sum(MyArray()), 'yes')
  286. @requires_array_function
  287. def test_sum_on_mock_array(self):
  288. # We need a proxy for mocks because __array_function__ is only looked
  289. # up in the class dict
  290. class ArrayProxy:
  291. def __init__(self, value):
  292. self.value = value
  293. def __array_function__(self, *args, **kwargs):
  294. return self.value.__array_function__(*args, **kwargs)
  295. def __array__(self, *args, **kwargs):
  296. return self.value.__array__(*args, **kwargs)
  297. proxy = ArrayProxy(mock.Mock(spec=ArrayProxy))
  298. proxy.value.__array_function__.return_value = 1
  299. result = np.sum(proxy)
  300. assert_equal(result, 1)
  301. proxy.value.__array_function__.assert_called_once_with(
  302. np.sum, (ArrayProxy,), (proxy,), {})
  303. proxy.value.__array__.assert_not_called()
  304. @requires_array_function
  305. def test_sum_forwarding_implementation(self):
  306. class MyArray(np.ndarray):
  307. def sum(self, axis, out):
  308. return 'summed'
  309. def __array_function__(self, func, types, args, kwargs):
  310. return super().__array_function__(func, types, args, kwargs)
  311. # note: the internal implementation of np.sum() calls the .sum() method
  312. array = np.array(1).view(MyArray)
  313. assert_equal(np.sum(array), 'summed')