codekingpro/portable-devtools

1import numbers
2import operator
3
4import numpy as np
5from numpy.testing import assert_, assert_equal, assert_raises
6
7# NOTE: This class should be kept as an exact copy of the example from the
8# docstring for NDArrayOperatorsMixin.
9
10class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin):
11    def __init__(self, value):
12        self.value = np.asarray(value)
13
14    # One might also consider adding the built-in list type to this
15    # list, to support operations like np.add(array_like, list)
16    _HANDLED_TYPES = (np.ndarray, numbers.Number)
17
18    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
19        out = kwargs.get('out', ())
20        for x in inputs + out:
21            # Only support operations with instances of _HANDLED_TYPES.
22            # Use ArrayLike instead of type(self) for isinstance to
23            # allow subclasses that don't override __array_ufunc__ to
24            # handle ArrayLike objects.
25            if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)):
26                return NotImplemented
27
28        # Defer to the implementation of the ufunc on unwrapped values.
29        inputs = tuple(x.value if isinstance(x, ArrayLike) else x
30                       for x in inputs)
31        if out:
32            kwargs['out'] = tuple(
33                x.value if isinstance(x, ArrayLike) else x
34                for x in out)
35        result = getattr(ufunc, method)(*inputs, **kwargs)
36
37        if type(result) is tuple:
38            # multiple return values
39            return tuple(type(self)(x) for x in result)
40        elif method == 'at':
41            # no return value
42            return None
43        else:
44            # one return value
45            return type(self)(result)
46
47    def __repr__(self):
48        return f'{type(self).__name__}({self.value!r})'
49
50
51def wrap_array_like(result):
52    if type(result) is tuple:
53        return tuple(ArrayLike(r) for r in result)
54    else:
55        return ArrayLike(result)
56
57
58def _assert_equal_type_and_value(result, expected, err_msg=None):
59    assert_equal(type(result), type(expected), err_msg=err_msg)
60    if isinstance(result, tuple):
61        assert_equal(len(result), len(expected), err_msg=err_msg)
62        for result_item, expected_item in zip(result, expected):
63            _assert_equal_type_and_value(result_item, expected_item, err_msg)
64    else:
65        assert_equal(result.value, expected.value, err_msg=err_msg)
66        assert_equal(getattr(result.value, 'dtype', None),
67                     getattr(expected.value, 'dtype', None), err_msg=err_msg)
68
69
70_ALL_BINARY_OPERATORS = [
71    operator.lt,
72    operator.le,
73    operator.eq,
74    operator.ne,
75    operator.gt,
76    operator.ge,
77    operator.add,
78    operator.sub,
79    operator.mul,
80    operator.truediv,
81    operator.floordiv,
82    operator.mod,
83    divmod,
84    pow,
85    operator.lshift,
86    operator.rshift,
87    operator.and_,
88    operator.xor,
89    operator.or_,
90]
91
92
93class TestNDArrayOperatorsMixin:
94
95    def test_array_like_add(self):
96
97        def check(result):
98            _assert_equal_type_and_value(result, ArrayLike(0))
99
100        check(ArrayLike(0) + 0)
101        check(0 + ArrayLike(0))
102
103        check(ArrayLike(0) + np.array(0))
104        check(np.array(0) + ArrayLike(0))
105
106        check(ArrayLike(np.array(0)) + 0)
107        check(0 + ArrayLike(np.array(0)))
108
109        check(ArrayLike(np.array(0)) + np.array(0))
110        check(np.array(0) + ArrayLike(np.array(0)))
111
112    def test_inplace(self):
113        array_like = ArrayLike(np.array([0]))
114        array_like += 1
115        _assert_equal_type_and_value(array_like, ArrayLike(np.array([1])))
116
117        array = np.array([0])
118        array += ArrayLike(1)
119        _assert_equal_type_and_value(array, ArrayLike(np.array([1])))
120
121    def test_opt_out(self):
122
123        class OptOut:
124            """Object that opts out of __array_ufunc__."""
125            __array_ufunc__ = None
126
127            def __add__(self, other):
128                return self
129
130            def __radd__(self, other):
131                return self
132
133        array_like = ArrayLike(1)
134        opt_out = OptOut()
135
136        # supported operations
137        assert_(array_like + opt_out is opt_out)
138        assert_(opt_out + array_like is opt_out)
139
140        # not supported
141        with assert_raises(TypeError):
142            # don't use the Python default, array_like = array_like + opt_out
143            array_like += opt_out
144        with assert_raises(TypeError):
145            array_like - opt_out
146        with assert_raises(TypeError):
147            opt_out - array_like
148
149    def test_subclass(self):
150
151        class SubArrayLike(ArrayLike):
152            """Should take precedence over ArrayLike."""
153
154        x = ArrayLike(0)
155        y = SubArrayLike(1)
156        _assert_equal_type_and_value(x + y, y)
157        _assert_equal_type_and_value(y + x, y)
158
159    def test_object(self):
160        x = ArrayLike(0)
161        obj = object()
162        with assert_raises(TypeError):
163            x + obj
164        with assert_raises(TypeError):
165            obj + x
166        with assert_raises(TypeError):
167            x += obj
168
169    def test_unary_methods(self):
170        array = np.array([-1, 0, 1, 2])
171        array_like = ArrayLike(array)
172        for op in [operator.neg,
173                   operator.pos,
174                   abs,
175                   operator.invert]:
176            _assert_equal_type_and_value(op(array_like), ArrayLike(op(array)))
177
178    def test_forward_binary_methods(self):
179        array = np.array([-1, 0, 1, 2])
180        array_like = ArrayLike(array)
181        for op in _ALL_BINARY_OPERATORS:
182            expected = wrap_array_like(op(array, 1))
183            actual = op(array_like, 1)
184            err_msg = f'failed for operator {op}'
185            _assert_equal_type_and_value(expected, actual, err_msg=err_msg)
186
187    def test_reflected_binary_methods(self):
188        for op in _ALL_BINARY_OPERATORS:
189            expected = wrap_array_like(op(2, 1))
190            actual = op(2, ArrayLike(1))
191            err_msg = f'failed for operator {op}'
192            _assert_equal_type_and_value(expected, actual, err_msg=err_msg)
193
194    def test_matmul(self):
195        array = np.array([1, 2], dtype=np.float64)
196        array_like = ArrayLike(array)
197        expected = ArrayLike(np.float64(5))
198        _assert_equal_type_and_value(expected, np.matmul(array_like, array))
199        _assert_equal_type_and_value(
200            expected, operator.matmul(array_like, array))
201        _assert_equal_type_and_value(
202            expected, operator.matmul(array, array_like))
203
204    def test_ufunc_at(self):
205        array = ArrayLike(np.array([1, 2, 3, 4]))
206        assert_(np.negative.at(array, np.array([0, 1])) is None)
207        _assert_equal_type_and_value(array, ArrayLike([-1, -2, 3, 4]))
208
209    def test_ufunc_two_outputs(self):
210        mantissa, exponent = np.frexp(2 ** -3)
211        expected = (ArrayLike(mantissa), ArrayLike(exponent))
212        _assert_equal_type_and_value(
213            np.frexp(ArrayLike(2 ** -3)), expected)
214        _assert_equal_type_and_value(
215            np.frexp(ArrayLike(np.array(2 ** -3))), expected)
216