123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216 |
- 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)
|