test_overrides.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583
  1. import inspect
  2. import sys
  3. import os
  4. import tempfile
  5. from io import StringIO
  6. from unittest import mock
  7. import numpy as np
  8. from numpy.testing import (
  9. assert_, assert_equal, assert_raises, assert_raises_regex)
  10. from numpy.core.overrides import (
  11. _get_implementing_args, array_function_dispatch,
  12. verify_matching_signatures, ARRAY_FUNCTION_ENABLED)
  13. from numpy.compat import pickle
  14. import pytest
  15. requires_array_function = pytest.mark.skipif(
  16. not ARRAY_FUNCTION_ENABLED,
  17. reason="__array_function__ dispatch not enabled.")
  18. def _return_not_implemented(self, *args, **kwargs):
  19. return NotImplemented
  20. # need to define this at the top level to test pickling
  21. @array_function_dispatch(lambda array: (array,))
  22. def dispatched_one_arg(array):
  23. """Docstring."""
  24. return 'original'
  25. @array_function_dispatch(lambda array1, array2: (array1, array2))
  26. def dispatched_two_arg(array1, array2):
  27. """Docstring."""
  28. return 'original'
  29. class TestGetImplementingArgs:
  30. def test_ndarray(self):
  31. array = np.array(1)
  32. args = _get_implementing_args([array])
  33. assert_equal(list(args), [array])
  34. args = _get_implementing_args([array, array])
  35. assert_equal(list(args), [array])
  36. args = _get_implementing_args([array, 1])
  37. assert_equal(list(args), [array])
  38. args = _get_implementing_args([1, array])
  39. assert_equal(list(args), [array])
  40. def test_ndarray_subclasses(self):
  41. class OverrideSub(np.ndarray):
  42. __array_function__ = _return_not_implemented
  43. class NoOverrideSub(np.ndarray):
  44. pass
  45. array = np.array(1).view(np.ndarray)
  46. override_sub = np.array(1).view(OverrideSub)
  47. no_override_sub = np.array(1).view(NoOverrideSub)
  48. args = _get_implementing_args([array, override_sub])
  49. assert_equal(list(args), [override_sub, array])
  50. args = _get_implementing_args([array, no_override_sub])
  51. assert_equal(list(args), [no_override_sub, array])
  52. args = _get_implementing_args(
  53. [override_sub, no_override_sub])
  54. assert_equal(list(args), [override_sub, no_override_sub])
  55. def test_ndarray_and_duck_array(self):
  56. class Other:
  57. __array_function__ = _return_not_implemented
  58. array = np.array(1)
  59. other = Other()
  60. args = _get_implementing_args([other, array])
  61. assert_equal(list(args), [other, array])
  62. args = _get_implementing_args([array, other])
  63. assert_equal(list(args), [array, other])
  64. def test_ndarray_subclass_and_duck_array(self):
  65. class OverrideSub(np.ndarray):
  66. __array_function__ = _return_not_implemented
  67. class Other:
  68. __array_function__ = _return_not_implemented
  69. array = np.array(1)
  70. subarray = np.array(1).view(OverrideSub)
  71. other = Other()
  72. assert_equal(_get_implementing_args([array, subarray, other]),
  73. [subarray, array, other])
  74. assert_equal(_get_implementing_args([array, other, subarray]),
  75. [subarray, array, other])
  76. def test_many_duck_arrays(self):
  77. class A:
  78. __array_function__ = _return_not_implemented
  79. class B(A):
  80. __array_function__ = _return_not_implemented
  81. class C(A):
  82. __array_function__ = _return_not_implemented
  83. class D:
  84. __array_function__ = _return_not_implemented
  85. a = A()
  86. b = B()
  87. c = C()
  88. d = D()
  89. assert_equal(_get_implementing_args([1]), [])
  90. assert_equal(_get_implementing_args([a]), [a])
  91. assert_equal(_get_implementing_args([a, 1]), [a])
  92. assert_equal(_get_implementing_args([a, a, a]), [a])
  93. assert_equal(_get_implementing_args([a, d, a]), [a, d])
  94. assert_equal(_get_implementing_args([a, b]), [b, a])
  95. assert_equal(_get_implementing_args([b, a]), [b, a])
  96. assert_equal(_get_implementing_args([a, b, c]), [b, c, a])
  97. assert_equal(_get_implementing_args([a, c, b]), [c, b, a])
  98. def test_too_many_duck_arrays(self):
  99. namespace = dict(__array_function__=_return_not_implemented)
  100. types = [type('A' + str(i), (object,), namespace) for i in range(33)]
  101. relevant_args = [t() for t in types]
  102. actual = _get_implementing_args(relevant_args[:32])
  103. assert_equal(actual, relevant_args[:32])
  104. with assert_raises_regex(TypeError, 'distinct argument types'):
  105. _get_implementing_args(relevant_args)
  106. class TestNDArrayArrayFunction:
  107. @requires_array_function
  108. def test_method(self):
  109. class Other:
  110. __array_function__ = _return_not_implemented
  111. class NoOverrideSub(np.ndarray):
  112. pass
  113. class OverrideSub(np.ndarray):
  114. __array_function__ = _return_not_implemented
  115. array = np.array([1])
  116. other = Other()
  117. no_override_sub = array.view(NoOverrideSub)
  118. override_sub = array.view(OverrideSub)
  119. result = array.__array_function__(func=dispatched_two_arg,
  120. types=(np.ndarray,),
  121. args=(array, 1.), kwargs={})
  122. assert_equal(result, 'original')
  123. result = array.__array_function__(func=dispatched_two_arg,
  124. types=(np.ndarray, Other),
  125. args=(array, other), kwargs={})
  126. assert_(result is NotImplemented)
  127. result = array.__array_function__(func=dispatched_two_arg,
  128. types=(np.ndarray, NoOverrideSub),
  129. args=(array, no_override_sub),
  130. kwargs={})
  131. assert_equal(result, 'original')
  132. result = array.__array_function__(func=dispatched_two_arg,
  133. types=(np.ndarray, OverrideSub),
  134. args=(array, override_sub),
  135. kwargs={})
  136. assert_equal(result, 'original')
  137. with assert_raises_regex(TypeError, 'no implementation found'):
  138. np.concatenate((array, other))
  139. expected = np.concatenate((array, array))
  140. result = np.concatenate((array, no_override_sub))
  141. assert_equal(result, expected.view(NoOverrideSub))
  142. result = np.concatenate((array, override_sub))
  143. assert_equal(result, expected.view(OverrideSub))
  144. def test_no_wrapper(self):
  145. # This shouldn't happen unless a user intentionally calls
  146. # __array_function__ with invalid arguments, but check that we raise
  147. # an appropriate error all the same.
  148. array = np.array(1)
  149. func = lambda x: x
  150. with assert_raises_regex(AttributeError, '_implementation'):
  151. array.__array_function__(func=func, types=(np.ndarray,),
  152. args=(array,), kwargs={})
  153. @requires_array_function
  154. class TestArrayFunctionDispatch:
  155. def test_pickle(self):
  156. for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
  157. roundtripped = pickle.loads(
  158. pickle.dumps(dispatched_one_arg, protocol=proto))
  159. assert_(roundtripped is dispatched_one_arg)
  160. def test_name_and_docstring(self):
  161. assert_equal(dispatched_one_arg.__name__, 'dispatched_one_arg')
  162. if sys.flags.optimize < 2:
  163. assert_equal(dispatched_one_arg.__doc__, 'Docstring.')
  164. def test_interface(self):
  165. class MyArray:
  166. def __array_function__(self, func, types, args, kwargs):
  167. return (self, func, types, args, kwargs)
  168. original = MyArray()
  169. (obj, func, types, args, kwargs) = dispatched_one_arg(original)
  170. assert_(obj is original)
  171. assert_(func is dispatched_one_arg)
  172. assert_equal(set(types), {MyArray})
  173. # assert_equal uses the overloaded np.iscomplexobj() internally
  174. assert_(args == (original,))
  175. assert_equal(kwargs, {})
  176. def test_not_implemented(self):
  177. class MyArray:
  178. def __array_function__(self, func, types, args, kwargs):
  179. return NotImplemented
  180. array = MyArray()
  181. with assert_raises_regex(TypeError, 'no implementation found'):
  182. dispatched_one_arg(array)
  183. @requires_array_function
  184. class TestVerifyMatchingSignatures:
  185. def test_verify_matching_signatures(self):
  186. verify_matching_signatures(lambda x: 0, lambda x: 0)
  187. verify_matching_signatures(lambda x=None: 0, lambda x=None: 0)
  188. verify_matching_signatures(lambda x=1: 0, lambda x=None: 0)
  189. with assert_raises(RuntimeError):
  190. verify_matching_signatures(lambda a: 0, lambda b: 0)
  191. with assert_raises(RuntimeError):
  192. verify_matching_signatures(lambda x: 0, lambda x=None: 0)
  193. with assert_raises(RuntimeError):
  194. verify_matching_signatures(lambda x=None: 0, lambda y=None: 0)
  195. with assert_raises(RuntimeError):
  196. verify_matching_signatures(lambda x=1: 0, lambda y=1: 0)
  197. def test_array_function_dispatch(self):
  198. with assert_raises(RuntimeError):
  199. @array_function_dispatch(lambda x: (x,))
  200. def f(y):
  201. pass
  202. # should not raise
  203. @array_function_dispatch(lambda x: (x,), verify=False)
  204. def f(y):
  205. pass
  206. def _new_duck_type_and_implements():
  207. """Create a duck array type and implements functions."""
  208. HANDLED_FUNCTIONS = {}
  209. class MyArray:
  210. def __array_function__(self, func, types, args, kwargs):
  211. if func not in HANDLED_FUNCTIONS:
  212. return NotImplemented
  213. if not all(issubclass(t, MyArray) for t in types):
  214. return NotImplemented
  215. return HANDLED_FUNCTIONS[func](*args, **kwargs)
  216. def implements(numpy_function):
  217. """Register an __array_function__ implementations."""
  218. def decorator(func):
  219. HANDLED_FUNCTIONS[numpy_function] = func
  220. return func
  221. return decorator
  222. return (MyArray, implements)
  223. @requires_array_function
  224. class TestArrayFunctionImplementation:
  225. def test_one_arg(self):
  226. MyArray, implements = _new_duck_type_and_implements()
  227. @implements(dispatched_one_arg)
  228. def _(array):
  229. return 'myarray'
  230. assert_equal(dispatched_one_arg(1), 'original')
  231. assert_equal(dispatched_one_arg(MyArray()), 'myarray')
  232. def test_optional_args(self):
  233. MyArray, implements = _new_duck_type_and_implements()
  234. @array_function_dispatch(lambda array, option=None: (array,))
  235. def func_with_option(array, option='default'):
  236. return option
  237. @implements(func_with_option)
  238. def my_array_func_with_option(array, new_option='myarray'):
  239. return new_option
  240. # we don't need to implement every option on __array_function__
  241. # implementations
  242. assert_equal(func_with_option(1), 'default')
  243. assert_equal(func_with_option(1, option='extra'), 'extra')
  244. assert_equal(func_with_option(MyArray()), 'myarray')
  245. with assert_raises(TypeError):
  246. func_with_option(MyArray(), option='extra')
  247. # but new options on implementations can't be used
  248. result = my_array_func_with_option(MyArray(), new_option='yes')
  249. assert_equal(result, 'yes')
  250. with assert_raises(TypeError):
  251. func_with_option(MyArray(), new_option='no')
  252. def test_not_implemented(self):
  253. MyArray, implements = _new_duck_type_and_implements()
  254. @array_function_dispatch(lambda array: (array,), module='my')
  255. def func(array):
  256. return array
  257. array = np.array(1)
  258. assert_(func(array) is array)
  259. assert_equal(func.__module__, 'my')
  260. with assert_raises_regex(
  261. TypeError, "no implementation found for 'my.func'"):
  262. func(MyArray())
  263. class TestNDArrayMethods:
  264. def test_repr(self):
  265. # gh-12162: should still be defined even if __array_function__ doesn't
  266. # implement np.array_repr()
  267. class MyArray(np.ndarray):
  268. def __array_function__(*args, **kwargs):
  269. return NotImplemented
  270. array = np.array(1).view(MyArray)
  271. assert_equal(repr(array), 'MyArray(1)')
  272. assert_equal(str(array), '1')
  273. class TestNumPyFunctions:
  274. def test_set_module(self):
  275. assert_equal(np.sum.__module__, 'numpy')
  276. assert_equal(np.char.equal.__module__, 'numpy.char')
  277. assert_equal(np.fft.fft.__module__, 'numpy.fft')
  278. assert_equal(np.linalg.solve.__module__, 'numpy.linalg')
  279. def test_inspect_sum(self):
  280. signature = inspect.signature(np.sum)
  281. assert_('axis' in signature.parameters)
  282. @requires_array_function
  283. def test_override_sum(self):
  284. MyArray, implements = _new_duck_type_and_implements()
  285. @implements(np.sum)
  286. def _(array):
  287. return 'yes'
  288. assert_equal(np.sum(MyArray()), 'yes')
  289. @requires_array_function
  290. def test_sum_on_mock_array(self):
  291. # We need a proxy for mocks because __array_function__ is only looked
  292. # up in the class dict
  293. class ArrayProxy:
  294. def __init__(self, value):
  295. self.value = value
  296. def __array_function__(self, *args, **kwargs):
  297. return self.value.__array_function__(*args, **kwargs)
  298. def __array__(self, *args, **kwargs):
  299. return self.value.__array__(*args, **kwargs)
  300. proxy = ArrayProxy(mock.Mock(spec=ArrayProxy))
  301. proxy.value.__array_function__.return_value = 1
  302. result = np.sum(proxy)
  303. assert_equal(result, 1)
  304. proxy.value.__array_function__.assert_called_once_with(
  305. np.sum, (ArrayProxy,), (proxy,), {})
  306. proxy.value.__array__.assert_not_called()
  307. @requires_array_function
  308. def test_sum_forwarding_implementation(self):
  309. class MyArray(np.ndarray):
  310. def sum(self, axis, out):
  311. return 'summed'
  312. def __array_function__(self, func, types, args, kwargs):
  313. return super().__array_function__(func, types, args, kwargs)
  314. # note: the internal implementation of np.sum() calls the .sum() method
  315. array = np.array(1).view(MyArray)
  316. assert_equal(np.sum(array), 'summed')
  317. class TestArrayLike:
  318. def setup(self):
  319. class MyArray():
  320. def __init__(self, function=None):
  321. self.function = function
  322. def __array_function__(self, func, types, args, kwargs):
  323. try:
  324. my_func = getattr(self, func.__name__)
  325. except AttributeError:
  326. return NotImplemented
  327. return my_func(*args, **kwargs)
  328. self.MyArray = MyArray
  329. class MyNoArrayFunctionArray():
  330. def __init__(self, function=None):
  331. self.function = function
  332. self.MyNoArrayFunctionArray = MyNoArrayFunctionArray
  333. def add_method(self, name, arr_class, enable_value_error=False):
  334. def _definition(*args, **kwargs):
  335. # Check that `like=` isn't propagated downstream
  336. assert 'like' not in kwargs
  337. if enable_value_error and 'value_error' in kwargs:
  338. raise ValueError
  339. return arr_class(getattr(arr_class, name))
  340. setattr(arr_class, name, _definition)
  341. def func_args(*args, **kwargs):
  342. return args, kwargs
  343. @requires_array_function
  344. def test_array_like_not_implemented(self):
  345. self.add_method('array', self.MyArray)
  346. ref = self.MyArray.array()
  347. with assert_raises_regex(TypeError, 'no implementation found'):
  348. array_like = np.asarray(1, like=ref)
  349. _array_tests = [
  350. ('array', *func_args((1,))),
  351. ('asarray', *func_args((1,))),
  352. ('asanyarray', *func_args((1,))),
  353. ('ascontiguousarray', *func_args((2, 3))),
  354. ('asfortranarray', *func_args((2, 3))),
  355. ('require', *func_args((np.arange(6).reshape(2, 3),),
  356. requirements=['A', 'F'])),
  357. ('empty', *func_args((1,))),
  358. ('full', *func_args((1,), 2)),
  359. ('ones', *func_args((1,))),
  360. ('zeros', *func_args((1,))),
  361. ('arange', *func_args(3)),
  362. ('frombuffer', *func_args(b'\x00' * 8, dtype=int)),
  363. ('fromiter', *func_args(range(3), dtype=int)),
  364. ('fromstring', *func_args('1,2', dtype=int, sep=',')),
  365. ('loadtxt', *func_args(lambda: StringIO('0 1\n2 3'))),
  366. ('genfromtxt', *func_args(lambda: StringIO(u'1,2.1'),
  367. dtype=[('int', 'i8'), ('float', 'f8')],
  368. delimiter=',')),
  369. ]
  370. @pytest.mark.parametrize('function, args, kwargs', _array_tests)
  371. @pytest.mark.parametrize('numpy_ref', [True, False])
  372. @requires_array_function
  373. def test_array_like(self, function, args, kwargs, numpy_ref):
  374. self.add_method('array', self.MyArray)
  375. self.add_method(function, self.MyArray)
  376. np_func = getattr(np, function)
  377. my_func = getattr(self.MyArray, function)
  378. if numpy_ref is True:
  379. ref = np.array(1)
  380. else:
  381. ref = self.MyArray.array()
  382. like_args = tuple(a() if callable(a) else a for a in args)
  383. array_like = np_func(*like_args, **kwargs, like=ref)
  384. if numpy_ref is True:
  385. assert type(array_like) is np.ndarray
  386. np_args = tuple(a() if callable(a) else a for a in args)
  387. np_arr = np_func(*np_args, **kwargs)
  388. # Special-case np.empty to ensure values match
  389. if function == "empty":
  390. np_arr.fill(1)
  391. array_like.fill(1)
  392. assert_equal(array_like, np_arr)
  393. else:
  394. assert type(array_like) is self.MyArray
  395. assert array_like.function is my_func
  396. @pytest.mark.parametrize('function, args, kwargs', _array_tests)
  397. @pytest.mark.parametrize('ref', [1, [1], "MyNoArrayFunctionArray"])
  398. @requires_array_function
  399. def test_no_array_function_like(self, function, args, kwargs, ref):
  400. self.add_method('array', self.MyNoArrayFunctionArray)
  401. self.add_method(function, self.MyNoArrayFunctionArray)
  402. np_func = getattr(np, function)
  403. # Instantiate ref if it's the MyNoArrayFunctionArray class
  404. if ref == "MyNoArrayFunctionArray":
  405. ref = self.MyNoArrayFunctionArray.array()
  406. like_args = tuple(a() if callable(a) else a for a in args)
  407. with assert_raises_regex(TypeError,
  408. 'The `like` argument must be an array-like that implements'):
  409. np_func(*like_args, **kwargs, like=ref)
  410. @pytest.mark.parametrize('numpy_ref', [True, False])
  411. def test_array_like_fromfile(self, numpy_ref):
  412. self.add_method('array', self.MyArray)
  413. self.add_method("fromfile", self.MyArray)
  414. if numpy_ref is True:
  415. ref = np.array(1)
  416. else:
  417. ref = self.MyArray.array()
  418. data = np.random.random(5)
  419. with tempfile.TemporaryDirectory() as tmpdir:
  420. fname = os.path.join(tmpdir, "testfile")
  421. data.tofile(fname)
  422. array_like = np.fromfile(fname, like=ref)
  423. if numpy_ref is True:
  424. assert type(array_like) is np.ndarray
  425. np_res = np.fromfile(fname, like=ref)
  426. assert_equal(np_res, data)
  427. assert_equal(array_like, np_res)
  428. else:
  429. assert type(array_like) is self.MyArray
  430. assert array_like.function is self.MyArray.fromfile
  431. @requires_array_function
  432. def test_exception_handling(self):
  433. self.add_method('array', self.MyArray, enable_value_error=True)
  434. ref = self.MyArray.array()
  435. with assert_raises(ValueError):
  436. np.array(1, value_error=True, like=ref)