test_subclassing.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460
  1. # pylint: disable-msg=W0611, W0612, W0511,R0201
  2. """Tests suite for MaskedArray & subclassing.
  3. :author: Pierre Gerard-Marchant
  4. :contact: pierregm_at_uga_dot_edu
  5. :version: $Id: test_subclassing.py 3473 2007-10-29 15:18:13Z jarrod.millman $
  6. """
  7. import numpy as np
  8. from numpy.lib.mixins import NDArrayOperatorsMixin
  9. from numpy.testing import assert_, assert_raises
  10. from numpy.ma.testutils import assert_equal
  11. from numpy.ma.core import (
  12. array, arange, masked, MaskedArray, masked_array, log, add, hypot,
  13. divide, asarray, asanyarray, nomask
  14. )
  15. # from numpy.ma.core import (
  16. def assert_startswith(a, b):
  17. # produces a better error message than assert_(a.startswith(b))
  18. assert_equal(a[:len(b)], b)
  19. class SubArray(np.ndarray):
  20. # Defines a generic np.ndarray subclass, that stores some metadata
  21. # in the dictionary `info`.
  22. def __new__(cls,arr,info={}):
  23. x = np.asanyarray(arr).view(cls)
  24. x.info = info.copy()
  25. return x
  26. def __array_finalize__(self, obj):
  27. super().__array_finalize__(obj)
  28. self.info = getattr(obj, 'info', {}).copy()
  29. return
  30. def __add__(self, other):
  31. result = super().__add__(other)
  32. result.info['added'] = result.info.get('added', 0) + 1
  33. return result
  34. def __iadd__(self, other):
  35. result = super().__iadd__(other)
  36. result.info['iadded'] = result.info.get('iadded', 0) + 1
  37. return result
  38. subarray = SubArray
  39. class SubMaskedArray(MaskedArray):
  40. """Pure subclass of MaskedArray, keeping some info on subclass."""
  41. def __new__(cls, info=None, **kwargs):
  42. obj = super().__new__(cls, **kwargs)
  43. obj._optinfo['info'] = info
  44. return obj
  45. class MSubArray(SubArray, MaskedArray):
  46. def __new__(cls, data, info={}, mask=nomask):
  47. subarr = SubArray(data, info)
  48. _data = MaskedArray.__new__(cls, data=subarr, mask=mask)
  49. _data.info = subarr.info
  50. return _data
  51. @property
  52. def _series(self):
  53. _view = self.view(MaskedArray)
  54. _view._sharedmask = False
  55. return _view
  56. msubarray = MSubArray
  57. # Also a subclass that overrides __str__, __repr__ and __setitem__, disallowing
  58. # setting to non-class values (and thus np.ma.core.masked_print_option)
  59. # and overrides __array_wrap__, updating the info dict, to check that this
  60. # doesn't get destroyed by MaskedArray._update_from. But this one also needs
  61. # its own iterator...
  62. class CSAIterator:
  63. """
  64. Flat iterator object that uses its own setter/getter
  65. (works around ndarray.flat not propagating subclass setters/getters
  66. see https://github.com/numpy/numpy/issues/4564)
  67. roughly following MaskedIterator
  68. """
  69. def __init__(self, a):
  70. self._original = a
  71. self._dataiter = a.view(np.ndarray).flat
  72. def __iter__(self):
  73. return self
  74. def __getitem__(self, indx):
  75. out = self._dataiter.__getitem__(indx)
  76. if not isinstance(out, np.ndarray):
  77. out = out.__array__()
  78. out = out.view(type(self._original))
  79. return out
  80. def __setitem__(self, index, value):
  81. self._dataiter[index] = self._original._validate_input(value)
  82. def __next__(self):
  83. return next(self._dataiter).__array__().view(type(self._original))
  84. class ComplicatedSubArray(SubArray):
  85. def __str__(self):
  86. return f'myprefix {self.view(SubArray)} mypostfix'
  87. def __repr__(self):
  88. # Return a repr that does not start with 'name('
  89. return f'<{self.__class__.__name__} {self}>'
  90. def _validate_input(self, value):
  91. if not isinstance(value, ComplicatedSubArray):
  92. raise ValueError("Can only set to MySubArray values")
  93. return value
  94. def __setitem__(self, item, value):
  95. # validation ensures direct assignment with ndarray or
  96. # masked_print_option will fail
  97. super().__setitem__(item, self._validate_input(value))
  98. def __getitem__(self, item):
  99. # ensure getter returns our own class also for scalars
  100. value = super().__getitem__(item)
  101. if not isinstance(value, np.ndarray): # scalar
  102. value = value.__array__().view(ComplicatedSubArray)
  103. return value
  104. @property
  105. def flat(self):
  106. return CSAIterator(self)
  107. @flat.setter
  108. def flat(self, value):
  109. y = self.ravel()
  110. y[:] = value
  111. def __array_wrap__(self, obj, context=None, return_scalar=False):
  112. obj = super().__array_wrap__(obj, context, return_scalar)
  113. if context is not None and context[0] is np.multiply:
  114. obj.info['multiplied'] = obj.info.get('multiplied', 0) + 1
  115. return obj
  116. class WrappedArray(NDArrayOperatorsMixin):
  117. """
  118. Wrapping a MaskedArray rather than subclassing to test that
  119. ufunc deferrals are commutative.
  120. See: https://github.com/numpy/numpy/issues/15200)
  121. """
  122. __slots__ = ('_array', 'attrs')
  123. __array_priority__ = 20
  124. def __init__(self, array, **attrs):
  125. self._array = array
  126. self.attrs = attrs
  127. def __repr__(self):
  128. return f"{self.__class__.__name__}(\n{self._array}\n{self.attrs}\n)"
  129. def __array__(self, dtype=None, copy=None):
  130. return np.asarray(self._array)
  131. def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
  132. if method == '__call__':
  133. inputs = [arg._array if isinstance(arg, self.__class__) else arg
  134. for arg in inputs]
  135. return self.__class__(ufunc(*inputs, **kwargs), **self.attrs)
  136. else:
  137. return NotImplemented
  138. class TestSubclassing:
  139. # Test suite for masked subclasses of ndarray.
  140. def setup_method(self):
  141. x = np.arange(5, dtype='float')
  142. mx = msubarray(x, mask=[0, 1, 0, 0, 0])
  143. self.data = (x, mx)
  144. def test_data_subclassing(self):
  145. # Tests whether the subclass is kept.
  146. x = np.arange(5)
  147. m = [0, 0, 1, 0, 0]
  148. xsub = SubArray(x)
  149. xmsub = masked_array(xsub, mask=m)
  150. assert_(isinstance(xmsub, MaskedArray))
  151. assert_equal(xmsub._data, xsub)
  152. assert_(isinstance(xmsub._data, SubArray))
  153. def test_maskedarray_subclassing(self):
  154. # Tests subclassing MaskedArray
  155. (x, mx) = self.data
  156. assert_(isinstance(mx._data, subarray))
  157. def test_masked_unary_operations(self):
  158. # Tests masked_unary_operation
  159. (x, mx) = self.data
  160. with np.errstate(divide='ignore'):
  161. assert_(isinstance(log(mx), msubarray))
  162. assert_equal(log(x), np.log(x))
  163. def test_masked_binary_operations(self):
  164. # Tests masked_binary_operation
  165. (x, mx) = self.data
  166. # Result should be a msubarray
  167. assert_(isinstance(add(mx, mx), msubarray))
  168. assert_(isinstance(add(mx, x), msubarray))
  169. # Result should work
  170. assert_equal(add(mx, x), mx+x)
  171. assert_(isinstance(add(mx, mx)._data, subarray))
  172. assert_(isinstance(add.outer(mx, mx), msubarray))
  173. assert_(isinstance(hypot(mx, mx), msubarray))
  174. assert_(isinstance(hypot(mx, x), msubarray))
  175. def test_masked_binary_operations2(self):
  176. # Tests domained_masked_binary_operation
  177. (x, mx) = self.data
  178. xmx = masked_array(mx.data.__array__(), mask=mx.mask)
  179. assert_(isinstance(divide(mx, mx), msubarray))
  180. assert_(isinstance(divide(mx, x), msubarray))
  181. assert_equal(divide(mx, mx), divide(xmx, xmx))
  182. def test_attributepropagation(self):
  183. x = array(arange(5), mask=[0]+[1]*4)
  184. my = masked_array(subarray(x))
  185. ym = msubarray(x)
  186. #
  187. z = (my+1)
  188. assert_(isinstance(z, MaskedArray))
  189. assert_(not isinstance(z, MSubArray))
  190. assert_(isinstance(z._data, SubArray))
  191. assert_equal(z._data.info, {})
  192. #
  193. z = (ym+1)
  194. assert_(isinstance(z, MaskedArray))
  195. assert_(isinstance(z, MSubArray))
  196. assert_(isinstance(z._data, SubArray))
  197. assert_(z._data.info['added'] > 0)
  198. # Test that inplace methods from data get used (gh-4617)
  199. ym += 1
  200. assert_(isinstance(ym, MaskedArray))
  201. assert_(isinstance(ym, MSubArray))
  202. assert_(isinstance(ym._data, SubArray))
  203. assert_(ym._data.info['iadded'] > 0)
  204. #
  205. ym._set_mask([1, 0, 0, 0, 1])
  206. assert_equal(ym._mask, [1, 0, 0, 0, 1])
  207. ym._series._set_mask([0, 0, 0, 0, 1])
  208. assert_equal(ym._mask, [0, 0, 0, 0, 1])
  209. #
  210. xsub = subarray(x, info={'name':'x'})
  211. mxsub = masked_array(xsub)
  212. assert_(hasattr(mxsub, 'info'))
  213. assert_equal(mxsub.info, xsub.info)
  214. def test_subclasspreservation(self):
  215. # Checks that masked_array(...,subok=True) preserves the class.
  216. x = np.arange(5)
  217. m = [0, 0, 1, 0, 0]
  218. xinfo = list(zip(x, m))
  219. xsub = MSubArray(x, mask=m, info={'xsub':xinfo})
  220. #
  221. mxsub = masked_array(xsub, subok=False)
  222. assert_(not isinstance(mxsub, MSubArray))
  223. assert_(isinstance(mxsub, MaskedArray))
  224. assert_equal(mxsub._mask, m)
  225. #
  226. mxsub = asarray(xsub)
  227. assert_(not isinstance(mxsub, MSubArray))
  228. assert_(isinstance(mxsub, MaskedArray))
  229. assert_equal(mxsub._mask, m)
  230. #
  231. mxsub = masked_array(xsub, subok=True)
  232. assert_(isinstance(mxsub, MSubArray))
  233. assert_equal(mxsub.info, xsub.info)
  234. assert_equal(mxsub._mask, xsub._mask)
  235. #
  236. mxsub = asanyarray(xsub)
  237. assert_(isinstance(mxsub, MSubArray))
  238. assert_equal(mxsub.info, xsub.info)
  239. assert_equal(mxsub._mask, m)
  240. def test_subclass_items(self):
  241. """test that getter and setter go via baseclass"""
  242. x = np.arange(5)
  243. xcsub = ComplicatedSubArray(x)
  244. mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
  245. # getter should return a ComplicatedSubArray, even for single item
  246. # first check we wrote ComplicatedSubArray correctly
  247. assert_(isinstance(xcsub[1], ComplicatedSubArray))
  248. assert_(isinstance(xcsub[1,...], ComplicatedSubArray))
  249. assert_(isinstance(xcsub[1:4], ComplicatedSubArray))
  250. # now that it propagates inside the MaskedArray
  251. assert_(isinstance(mxcsub[1], ComplicatedSubArray))
  252. assert_(isinstance(mxcsub[1,...].data, ComplicatedSubArray))
  253. assert_(mxcsub[0] is masked)
  254. assert_(isinstance(mxcsub[0,...].data, ComplicatedSubArray))
  255. assert_(isinstance(mxcsub[1:4].data, ComplicatedSubArray))
  256. # also for flattened version (which goes via MaskedIterator)
  257. assert_(isinstance(mxcsub.flat[1].data, ComplicatedSubArray))
  258. assert_(mxcsub.flat[0] is masked)
  259. assert_(isinstance(mxcsub.flat[1:4].base, ComplicatedSubArray))
  260. # setter should only work with ComplicatedSubArray input
  261. # first check we wrote ComplicatedSubArray correctly
  262. assert_raises(ValueError, xcsub.__setitem__, 1, x[4])
  263. # now that it propagates inside the MaskedArray
  264. assert_raises(ValueError, mxcsub.__setitem__, 1, x[4])
  265. assert_raises(ValueError, mxcsub.__setitem__, slice(1, 4), x[1:4])
  266. mxcsub[1] = xcsub[4]
  267. mxcsub[1:4] = xcsub[1:4]
  268. # also for flattened version (which goes via MaskedIterator)
  269. assert_raises(ValueError, mxcsub.flat.__setitem__, 1, x[4])
  270. assert_raises(ValueError, mxcsub.flat.__setitem__, slice(1, 4), x[1:4])
  271. mxcsub.flat[1] = xcsub[4]
  272. mxcsub.flat[1:4] = xcsub[1:4]
  273. def test_subclass_nomask_items(self):
  274. x = np.arange(5)
  275. xcsub = ComplicatedSubArray(x)
  276. mxcsub_nomask = masked_array(xcsub)
  277. assert_(isinstance(mxcsub_nomask[1,...].data, ComplicatedSubArray))
  278. assert_(isinstance(mxcsub_nomask[0,...].data, ComplicatedSubArray))
  279. assert_(isinstance(mxcsub_nomask[1], ComplicatedSubArray))
  280. assert_(isinstance(mxcsub_nomask[0], ComplicatedSubArray))
  281. def test_subclass_repr(self):
  282. """test that repr uses the name of the subclass
  283. and 'array' for np.ndarray"""
  284. x = np.arange(5)
  285. mx = masked_array(x, mask=[True, False, True, False, False])
  286. assert_startswith(repr(mx), 'masked_array')
  287. xsub = SubArray(x)
  288. mxsub = masked_array(xsub, mask=[True, False, True, False, False])
  289. assert_startswith(repr(mxsub),
  290. f'masked_{SubArray.__name__}(data=[--, 1, --, 3, 4]')
  291. def test_subclass_str(self):
  292. """test str with subclass that has overridden str, setitem"""
  293. # first without override
  294. x = np.arange(5)
  295. xsub = SubArray(x)
  296. mxsub = masked_array(xsub, mask=[True, False, True, False, False])
  297. assert_equal(str(mxsub), '[-- 1 -- 3 4]')
  298. xcsub = ComplicatedSubArray(x)
  299. assert_raises(ValueError, xcsub.__setitem__, 0,
  300. np.ma.core.masked_print_option)
  301. mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
  302. assert_equal(str(mxcsub), 'myprefix [-- 1 -- 3 4] mypostfix')
  303. def test_pure_subclass_info_preservation(self):
  304. # Test that ufuncs and methods conserve extra information consistently;
  305. # see gh-7122.
  306. arr1 = SubMaskedArray('test', data=[1,2,3,4,5,6])
  307. arr2 = SubMaskedArray(data=[0,1,2,3,4,5])
  308. diff1 = np.subtract(arr1, arr2)
  309. assert_('info' in diff1._optinfo)
  310. assert_(diff1._optinfo['info'] == 'test')
  311. diff2 = arr1 - arr2
  312. assert_('info' in diff2._optinfo)
  313. assert_(diff2._optinfo['info'] == 'test')
  314. class ArrayNoInheritance:
  315. """Quantity-like class that does not inherit from ndarray"""
  316. def __init__(self, data, units):
  317. self.magnitude = data
  318. self.units = units
  319. def __getattr__(self, attr):
  320. return getattr(self.magnitude, attr)
  321. def test_array_no_inheritance():
  322. data_masked = np.ma.array([1, 2, 3], mask=[True, False, True])
  323. data_masked_units = ArrayNoInheritance(data_masked, 'meters')
  324. # Get the masked representation of the Quantity-like class
  325. new_array = np.ma.array(data_masked_units)
  326. assert_equal(data_masked.data, new_array.data)
  327. assert_equal(data_masked.mask, new_array.mask)
  328. # Test sharing the mask
  329. data_masked.mask = [True, False, False]
  330. assert_equal(data_masked.mask, new_array.mask)
  331. assert_(new_array.sharedmask)
  332. # Get the masked representation of the Quantity-like class
  333. new_array = np.ma.array(data_masked_units, copy=True)
  334. assert_equal(data_masked.data, new_array.data)
  335. assert_equal(data_masked.mask, new_array.mask)
  336. # Test that the mask is not shared when copy=True
  337. data_masked.mask = [True, False, True]
  338. assert_equal([True, False, False], new_array.mask)
  339. assert_(not new_array.sharedmask)
  340. # Get the masked representation of the Quantity-like class
  341. new_array = np.ma.array(data_masked_units, keep_mask=False)
  342. assert_equal(data_masked.data, new_array.data)
  343. # The change did not affect the original mask
  344. assert_equal(data_masked.mask, [True, False, True])
  345. # Test that the mask is False and not shared when keep_mask=False
  346. assert_(not new_array.mask)
  347. assert_(not new_array.sharedmask)
  348. class TestClassWrapping:
  349. # Test suite for classes that wrap MaskedArrays
  350. def setup_method(self):
  351. m = np.ma.masked_array([1, 3, 5], mask=[False, True, False])
  352. wm = WrappedArray(m)
  353. self.data = (m, wm)
  354. def test_masked_unary_operations(self):
  355. # Tests masked_unary_operation
  356. (m, wm) = self.data
  357. with np.errstate(divide='ignore'):
  358. assert_(isinstance(np.log(wm), WrappedArray))
  359. def test_masked_binary_operations(self):
  360. # Tests masked_binary_operation
  361. (m, wm) = self.data
  362. # Result should be a WrappedArray
  363. assert_(isinstance(np.add(wm, wm), WrappedArray))
  364. assert_(isinstance(np.add(m, wm), WrappedArray))
  365. assert_(isinstance(np.add(wm, m), WrappedArray))
  366. # add and '+' should call the same ufunc
  367. assert_equal(np.add(m, wm), m + wm)
  368. assert_(isinstance(np.hypot(m, wm), WrappedArray))
  369. assert_(isinstance(np.hypot(wm, m), WrappedArray))
  370. # Test domained binary operations
  371. assert_(isinstance(np.divide(wm, m), WrappedArray))
  372. assert_(isinstance(np.divide(m, wm), WrappedArray))
  373. assert_equal(np.divide(wm, m) * m, np.divide(m, m) * wm)
  374. # Test broadcasting
  375. m2 = np.stack([m, m])
  376. assert_(isinstance(np.divide(wm, m2), WrappedArray))
  377. assert_(isinstance(np.divide(m2, wm), WrappedArray))
  378. assert_equal(np.divide(m2, wm), np.divide(wm, m2))
  379. def test_mixins_have_slots(self):
  380. mixin = NDArrayOperatorsMixin()
  381. # Should raise an error
  382. assert_raises(AttributeError, mixin.__setattr__, "not_a_real_attr", 1)
  383. m = np.ma.masked_array([1, 3, 5], mask=[False, True, False])
  384. wm = WrappedArray(m)
  385. assert_raises(AttributeError, wm.__setattr__, "not_an_attr", 2)