| import numbers
|
| import operator
|
|
|
| import numpy as np
|
| from numpy.testing import assert_, assert_equal, assert_raises
|
|
|
|
|
|
|
|
|
|
|
| class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin):
|
| def __init__(self, value):
|
| self.value = np.asarray(value)
|
|
|
|
|
|
|
| _HANDLED_TYPES = (np.ndarray, numbers.Number)
|
|
|
| def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
|
| out = kwargs.get('out', ())
|
| for x in inputs + out:
|
|
|
|
|
|
|
|
|
| if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)):
|
| return NotImplemented
|
|
|
|
|
| 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:
|
|
|
| return tuple(type(self)(x) for x in result)
|
| elif method == 'at':
|
|
|
| return None
|
| else:
|
|
|
| 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()
|
|
|
|
|
| assert_(array_like + opt_out is opt_out)
|
| assert_(opt_out + array_like is opt_out)
|
|
|
|
|
| with assert_raises(TypeError):
|
|
|
| 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)
|
|
|