JezK
Edit File: test_mixins.py
import numbers import operator import numpy as np from numpy.testing import assert_, assert_equal, assert_raises # NOTE: This class should be kept as an exact copy of the example from the # docstring for NDArrayOperatorsMixin. class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin): def __init__(self, value): self.value = np.asarray(value) # One might also consider adding the built-in list type to this # list, to support operations like np.add(array_like, list) _HANDLED_TYPES = (np.ndarray, numbers.Number) def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): out = kwargs.get('out', ()) for x in inputs + out: # Only support operations with instances of _HANDLED_TYPES. # Use ArrayLike instead of type(self) for isinstance to # allow subclasses that don't override __array_ufunc__ to # handle ArrayLike objects. if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)): return NotImplemented # Defer to the implementation of the ufunc on unwrapped values. inputs = tuple(x.value if isinstance(x, ArrayLike) else x for x in inputs) if out: kwargs['out'] = tuple( x.value if isinstance(x, ArrayLike) else x for x in out) result = getattr(ufunc, method)(*inputs, **kwargs) if type(result) is tuple: # multiple return values return tuple(type(self)(x) for x in result) elif method == 'at': # no return value return None else: # one return value return type(self)(result) def __repr__(self): return '%s(%r)' % (type(self).__name__, self.value) def wrap_array_like(result): if type(result) is tuple: return tuple(ArrayLike(r) for r in result) else: return ArrayLike(result) def _assert_equal_type_and_value(result, expected, err_msg=None): assert_equal(type(result), type(expected), err_msg=err_msg) if isinstance(result, tuple): assert_equal(len(result), len(expected), err_msg=err_msg) for result_item, expected_item in zip(result, expected): _assert_equal_type_and_value(result_item, expected_item, err_msg) else: assert_equal(result.value, expected.value, err_msg=err_msg) assert_equal(getattr(result.value, 'dtype', None), getattr(expected.value, 'dtype', None), err_msg=err_msg) _ALL_BINARY_OPERATORS = [ operator.lt, operator.le, operator.eq, operator.ne, operator.gt, operator.ge, operator.add, operator.sub, operator.mul, operator.truediv, operator.floordiv, operator.mod, divmod, pow, operator.lshift, operator.rshift, operator.and_, operator.xor, operator.or_, ] class TestNDArrayOperatorsMixin: def test_array_like_add(self): def check(result): _assert_equal_type_and_value(result, ArrayLike(0)) check(ArrayLike(0) + 0) check(0 + ArrayLike(0)) check(ArrayLike(0) + np.array(0)) check(np.array(0) + ArrayLike(0)) check(ArrayLike(np.array(0)) + 0) check(0 + ArrayLike(np.array(0))) check(ArrayLike(np.array(0)) + np.array(0)) check(np.array(0) + ArrayLike(np.array(0))) def test_inplace(self): array_like = ArrayLike(np.array([0])) array_like += 1 _assert_equal_type_and_value(array_like, ArrayLike(np.array([1]))) array = np.array([0]) array += ArrayLike(1) _assert_equal_type_and_value(array, ArrayLike(np.array([1]))) def test_opt_out(self): class OptOut: """Object that opts out of __array_ufunc__.""" __array_ufunc__ = None def __add__(self, other): return self def __radd__(self, other): return self array_like = ArrayLike(1) opt_out = OptOut() # supported operations assert_(array_like + opt_out is opt_out) assert_(opt_out + array_like is opt_out) # not supported with assert_raises(TypeError): # don't use the Python default, array_like = array_like + opt_out array_like += opt_out with assert_raises(TypeError): array_like - opt_out with assert_raises(TypeError): opt_out - array_like def test_subclass(self): class SubArrayLike(ArrayLike): """Should take precedence over ArrayLike.""" x = ArrayLike(0) y = SubArrayLike(1) _assert_equal_type_and_value(x + y, y) _assert_equal_type_and_value(y + x, y) def test_object(self): x = ArrayLike(0) obj = object() with assert_raises(TypeError): x + obj with assert_raises(TypeError): obj + x with assert_raises(TypeError): x += obj def test_unary_methods(self): array = np.array([-1, 0, 1, 2]) array_like = ArrayLike(array) for op in [operator.neg, operator.pos, abs, operator.invert]: _assert_equal_type_and_value(op(array_like), ArrayLike(op(array))) def test_forward_binary_methods(self): array = np.array([-1, 0, 1, 2]) array_like = ArrayLike(array) for op in _ALL_BINARY_OPERATORS: expected = wrap_array_like(op(array, 1)) actual = op(array_like, 1) err_msg = 'failed for operator {}'.format(op) _assert_equal_type_and_value(expected, actual, err_msg=err_msg) def test_reflected_binary_methods(self): for op in _ALL_BINARY_OPERATORS: expected = wrap_array_like(op(2, 1)) actual = op(2, ArrayLike(1)) err_msg = 'failed for operator {}'.format(op) _assert_equal_type_and_value(expected, actual, err_msg=err_msg) def test_matmul(self): array = np.array([1, 2], dtype=np.float64) array_like = ArrayLike(array) expected = ArrayLike(np.float64(5)) _assert_equal_type_and_value(expected, np.matmul(array_like, array)) _assert_equal_type_and_value( expected, operator.matmul(array_like, array)) _assert_equal_type_and_value( expected, operator.matmul(array, array_like)) def test_ufunc_at(self): array = ArrayLike(np.array([1, 2, 3, 4])) assert_(np.negative.at(array, np.array([0, 1])) is None) _assert_equal_type_and_value(array, ArrayLike([-1, -2, 3, 4])) def test_ufunc_two_outputs(self): mantissa, exponent = np.frexp(2 ** -3) expected = (ArrayLike(mantissa), ArrayLike(exponent)) _assert_equal_type_and_value( np.frexp(ArrayLike(2 ** -3)), expected) _assert_equal_type_and_value( np.frexp(ArrayLike(np.array(2 ** -3))), expected)