test_subclassing.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347
  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.testing import assert_, assert_raises
  9. from numpy.ma.testutils import assert_equal
  10. from numpy.ma.core import (
  11. array, arange, masked, MaskedArray, masked_array, log, add, hypot,
  12. divide, asarray, asanyarray, nomask
  13. )
  14. # from numpy.ma.core import (
  15. def assert_startswith(a, b):
  16. # produces a better error message than assert_(a.startswith(b))
  17. assert_equal(a[:len(b)], b)
  18. class SubArray(np.ndarray):
  19. # Defines a generic np.ndarray subclass, that stores some metadata
  20. # in the dictionary `info`.
  21. def __new__(cls,arr,info={}):
  22. x = np.asanyarray(arr).view(cls)
  23. x.info = info.copy()
  24. return x
  25. def __array_finalize__(self, obj):
  26. if callable(getattr(super(SubArray, self),
  27. '__array_finalize__', None)):
  28. super(SubArray, self).__array_finalize__(obj)
  29. self.info = getattr(obj, 'info', {}).copy()
  30. return
  31. def __add__(self, other):
  32. result = super(SubArray, self).__add__(other)
  33. result.info['added'] = result.info.get('added', 0) + 1
  34. return result
  35. def __iadd__(self, other):
  36. result = super(SubArray, self).__iadd__(other)
  37. result.info['iadded'] = result.info.get('iadded', 0) + 1
  38. return result
  39. subarray = SubArray
  40. class SubMaskedArray(MaskedArray):
  41. """Pure subclass of MaskedArray, keeping some info on subclass."""
  42. def __new__(cls, info=None, **kwargs):
  43. obj = super(SubMaskedArray, cls).__new__(cls, **kwargs)
  44. obj._optinfo['info'] = info
  45. return obj
  46. class MSubArray(SubArray, MaskedArray):
  47. def __new__(cls, data, info={}, mask=nomask):
  48. subarr = SubArray(data, info)
  49. _data = MaskedArray.__new__(cls, data=subarr, mask=mask)
  50. _data.info = subarr.info
  51. return _data
  52. @property
  53. def _series(self):
  54. _view = self.view(MaskedArray)
  55. _view._sharedmask = False
  56. return _view
  57. msubarray = MSubArray
  58. # Also a subclass that overrides __str__, __repr__ and __setitem__, disallowing
  59. # setting to non-class values (and thus np.ma.core.masked_print_option)
  60. # and overrides __array_wrap__, updating the info dict, to check that this
  61. # doesn't get destroyed by MaskedArray._update_from. But this one also needs
  62. # its own iterator...
  63. class CSAIterator:
  64. """
  65. Flat iterator object that uses its own setter/getter
  66. (works around ndarray.flat not propagating subclass setters/getters
  67. see https://github.com/numpy/numpy/issues/4564)
  68. roughly following MaskedIterator
  69. """
  70. def __init__(self, a):
  71. self._original = a
  72. self._dataiter = a.view(np.ndarray).flat
  73. def __iter__(self):
  74. return self
  75. def __getitem__(self, indx):
  76. out = self._dataiter.__getitem__(indx)
  77. if not isinstance(out, np.ndarray):
  78. out = out.__array__()
  79. out = out.view(type(self._original))
  80. return out
  81. def __setitem__(self, index, value):
  82. self._dataiter[index] = self._original._validate_input(value)
  83. def __next__(self):
  84. return next(self._dataiter).__array__().view(type(self._original))
  85. class ComplicatedSubArray(SubArray):
  86. def __str__(self):
  87. return f'myprefix {self.view(SubArray)} mypostfix'
  88. def __repr__(self):
  89. # Return a repr that does not start with 'name('
  90. return f'<{self.__class__.__name__} {self}>'
  91. def _validate_input(self, value):
  92. if not isinstance(value, ComplicatedSubArray):
  93. raise ValueError("Can only set to MySubArray values")
  94. return value
  95. def __setitem__(self, item, value):
  96. # validation ensures direct assignment with ndarray or
  97. # masked_print_option will fail
  98. super(ComplicatedSubArray, self).__setitem__(
  99. item, self._validate_input(value))
  100. def __getitem__(self, item):
  101. # ensure getter returns our own class also for scalars
  102. value = super(ComplicatedSubArray, self).__getitem__(item)
  103. if not isinstance(value, np.ndarray): # scalar
  104. value = value.__array__().view(ComplicatedSubArray)
  105. return value
  106. @property
  107. def flat(self):
  108. return CSAIterator(self)
  109. @flat.setter
  110. def flat(self, value):
  111. y = self.ravel()
  112. y[:] = value
  113. def __array_wrap__(self, obj, context=None):
  114. obj = super(ComplicatedSubArray, self).__array_wrap__(obj, context)
  115. if context is not None and context[0] is np.multiply:
  116. obj.info['multiplied'] = obj.info.get('multiplied', 0) + 1
  117. return obj
  118. class TestSubclassing:
  119. # Test suite for masked subclasses of ndarray.
  120. def setup(self):
  121. x = np.arange(5, dtype='float')
  122. mx = msubarray(x, mask=[0, 1, 0, 0, 0])
  123. self.data = (x, mx)
  124. def test_data_subclassing(self):
  125. # Tests whether the subclass is kept.
  126. x = np.arange(5)
  127. m = [0, 0, 1, 0, 0]
  128. xsub = SubArray(x)
  129. xmsub = masked_array(xsub, mask=m)
  130. assert_(isinstance(xmsub, MaskedArray))
  131. assert_equal(xmsub._data, xsub)
  132. assert_(isinstance(xmsub._data, SubArray))
  133. def test_maskedarray_subclassing(self):
  134. # Tests subclassing MaskedArray
  135. (x, mx) = self.data
  136. assert_(isinstance(mx._data, subarray))
  137. def test_masked_unary_operations(self):
  138. # Tests masked_unary_operation
  139. (x, mx) = self.data
  140. with np.errstate(divide='ignore'):
  141. assert_(isinstance(log(mx), msubarray))
  142. assert_equal(log(x), np.log(x))
  143. def test_masked_binary_operations(self):
  144. # Tests masked_binary_operation
  145. (x, mx) = self.data
  146. # Result should be a msubarray
  147. assert_(isinstance(add(mx, mx), msubarray))
  148. assert_(isinstance(add(mx, x), msubarray))
  149. # Result should work
  150. assert_equal(add(mx, x), mx+x)
  151. assert_(isinstance(add(mx, mx)._data, subarray))
  152. assert_(isinstance(add.outer(mx, mx), msubarray))
  153. assert_(isinstance(hypot(mx, mx), msubarray))
  154. assert_(isinstance(hypot(mx, x), msubarray))
  155. def test_masked_binary_operations2(self):
  156. # Tests domained_masked_binary_operation
  157. (x, mx) = self.data
  158. xmx = masked_array(mx.data.__array__(), mask=mx.mask)
  159. assert_(isinstance(divide(mx, mx), msubarray))
  160. assert_(isinstance(divide(mx, x), msubarray))
  161. assert_equal(divide(mx, mx), divide(xmx, xmx))
  162. def test_attributepropagation(self):
  163. x = array(arange(5), mask=[0]+[1]*4)
  164. my = masked_array(subarray(x))
  165. ym = msubarray(x)
  166. #
  167. z = (my+1)
  168. assert_(isinstance(z, MaskedArray))
  169. assert_(not isinstance(z, MSubArray))
  170. assert_(isinstance(z._data, SubArray))
  171. assert_equal(z._data.info, {})
  172. #
  173. z = (ym+1)
  174. assert_(isinstance(z, MaskedArray))
  175. assert_(isinstance(z, MSubArray))
  176. assert_(isinstance(z._data, SubArray))
  177. assert_(z._data.info['added'] > 0)
  178. # Test that inplace methods from data get used (gh-4617)
  179. ym += 1
  180. assert_(isinstance(ym, MaskedArray))
  181. assert_(isinstance(ym, MSubArray))
  182. assert_(isinstance(ym._data, SubArray))
  183. assert_(ym._data.info['iadded'] > 0)
  184. #
  185. ym._set_mask([1, 0, 0, 0, 1])
  186. assert_equal(ym._mask, [1, 0, 0, 0, 1])
  187. ym._series._set_mask([0, 0, 0, 0, 1])
  188. assert_equal(ym._mask, [0, 0, 0, 0, 1])
  189. #
  190. xsub = subarray(x, info={'name':'x'})
  191. mxsub = masked_array(xsub)
  192. assert_(hasattr(mxsub, 'info'))
  193. assert_equal(mxsub.info, xsub.info)
  194. def test_subclasspreservation(self):
  195. # Checks that masked_array(...,subok=True) preserves the class.
  196. x = np.arange(5)
  197. m = [0, 0, 1, 0, 0]
  198. xinfo = [(i, j) for (i, j) in zip(x, m)]
  199. xsub = MSubArray(x, mask=m, info={'xsub':xinfo})
  200. #
  201. mxsub = masked_array(xsub, subok=False)
  202. assert_(not isinstance(mxsub, MSubArray))
  203. assert_(isinstance(mxsub, MaskedArray))
  204. assert_equal(mxsub._mask, m)
  205. #
  206. mxsub = asarray(xsub)
  207. assert_(not isinstance(mxsub, MSubArray))
  208. assert_(isinstance(mxsub, MaskedArray))
  209. assert_equal(mxsub._mask, m)
  210. #
  211. mxsub = masked_array(xsub, subok=True)
  212. assert_(isinstance(mxsub, MSubArray))
  213. assert_equal(mxsub.info, xsub.info)
  214. assert_equal(mxsub._mask, xsub._mask)
  215. #
  216. mxsub = asanyarray(xsub)
  217. assert_(isinstance(mxsub, MSubArray))
  218. assert_equal(mxsub.info, xsub.info)
  219. assert_equal(mxsub._mask, m)
  220. def test_subclass_items(self):
  221. """test that getter and setter go via baseclass"""
  222. x = np.arange(5)
  223. xcsub = ComplicatedSubArray(x)
  224. mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
  225. # getter should return a ComplicatedSubArray, even for single item
  226. # first check we wrote ComplicatedSubArray correctly
  227. assert_(isinstance(xcsub[1], ComplicatedSubArray))
  228. assert_(isinstance(xcsub[1,...], ComplicatedSubArray))
  229. assert_(isinstance(xcsub[1:4], ComplicatedSubArray))
  230. # now that it propagates inside the MaskedArray
  231. assert_(isinstance(mxcsub[1], ComplicatedSubArray))
  232. assert_(isinstance(mxcsub[1,...].data, ComplicatedSubArray))
  233. assert_(mxcsub[0] is masked)
  234. assert_(isinstance(mxcsub[0,...].data, ComplicatedSubArray))
  235. assert_(isinstance(mxcsub[1:4].data, ComplicatedSubArray))
  236. # also for flattened version (which goes via MaskedIterator)
  237. assert_(isinstance(mxcsub.flat[1].data, ComplicatedSubArray))
  238. assert_(mxcsub.flat[0] is masked)
  239. assert_(isinstance(mxcsub.flat[1:4].base, ComplicatedSubArray))
  240. # setter should only work with ComplicatedSubArray input
  241. # first check we wrote ComplicatedSubArray correctly
  242. assert_raises(ValueError, xcsub.__setitem__, 1, x[4])
  243. # now that it propagates inside the MaskedArray
  244. assert_raises(ValueError, mxcsub.__setitem__, 1, x[4])
  245. assert_raises(ValueError, mxcsub.__setitem__, slice(1, 4), x[1:4])
  246. mxcsub[1] = xcsub[4]
  247. mxcsub[1:4] = xcsub[1:4]
  248. # also for flattened version (which goes via MaskedIterator)
  249. assert_raises(ValueError, mxcsub.flat.__setitem__, 1, x[4])
  250. assert_raises(ValueError, mxcsub.flat.__setitem__, slice(1, 4), x[1:4])
  251. mxcsub.flat[1] = xcsub[4]
  252. mxcsub.flat[1:4] = xcsub[1:4]
  253. def test_subclass_nomask_items(self):
  254. x = np.arange(5)
  255. xcsub = ComplicatedSubArray(x)
  256. mxcsub_nomask = masked_array(xcsub)
  257. assert_(isinstance(mxcsub_nomask[1,...].data, ComplicatedSubArray))
  258. assert_(isinstance(mxcsub_nomask[0,...].data, ComplicatedSubArray))
  259. assert_(isinstance(mxcsub_nomask[1], ComplicatedSubArray))
  260. assert_(isinstance(mxcsub_nomask[0], ComplicatedSubArray))
  261. def test_subclass_repr(self):
  262. """test that repr uses the name of the subclass
  263. and 'array' for np.ndarray"""
  264. x = np.arange(5)
  265. mx = masked_array(x, mask=[True, False, True, False, False])
  266. assert_startswith(repr(mx), 'masked_array')
  267. xsub = SubArray(x)
  268. mxsub = masked_array(xsub, mask=[True, False, True, False, False])
  269. assert_startswith(repr(mxsub),
  270. f'masked_{SubArray.__name__}(data=[--, 1, --, 3, 4]')
  271. def test_subclass_str(self):
  272. """test str with subclass that has overridden str, setitem"""
  273. # first without override
  274. x = np.arange(5)
  275. xsub = SubArray(x)
  276. mxsub = masked_array(xsub, mask=[True, False, True, False, False])
  277. assert_equal(str(mxsub), '[-- 1 -- 3 4]')
  278. xcsub = ComplicatedSubArray(x)
  279. assert_raises(ValueError, xcsub.__setitem__, 0,
  280. np.ma.core.masked_print_option)
  281. mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
  282. assert_equal(str(mxcsub), 'myprefix [-- 1 -- 3 4] mypostfix')
  283. def test_pure_subclass_info_preservation(self):
  284. # Test that ufuncs and methods conserve extra information consistently;
  285. # see gh-7122.
  286. arr1 = SubMaskedArray('test', data=[1,2,3,4,5,6])
  287. arr2 = SubMaskedArray(data=[0,1,2,3,4,5])
  288. diff1 = np.subtract(arr1, arr2)
  289. assert_('info' in diff1._optinfo)
  290. assert_(diff1._optinfo['info'] == 'test')
  291. diff2 = arr1 - arr2
  292. assert_('info' in diff2._optinfo)
  293. assert_(diff2._optinfo['info'] == 'test')