ENH: Add NDArrayOperatorsMixin mixin class.

This mixin class provides an easy way to implement arithmetic operators
that defer to __array_ufunc__ like numpy.ndarray in non-ndarray
subclasses.

3 files changed, 358 insertions, 0 deletions

diff --git a/numpy/lib/__init__.py b/numpy/lib/__init__.py
index 1d65db55e..4cdb76b20 100644
--- a/numpy/lib/__init__.py
+++ b/numpy/lib/__init__.py
@@ -8,6 +8,7 @@ from numpy.version import version as __version__
 from .type_check import *
 from .index_tricks import *
 from .function_base import *
+from .mixins import *
 from .nanfunctions import *
 from .shape_base import *
 from .stride_tricks import *
@@ -29,6 +30,7 @@ __all__ = ['emath', 'math']
 __all__ += type_check.__all__
 __all__ += index_tricks.__all__
 __all__ += function_base.__all__
+__all__ += mixins.__all__
 __all__ += shape_base.__all__
 __all__ += stride_tricks.__all__
 __all__ += twodim_base.__all__
diff --git a/numpy/lib/mixins.py b/numpy/lib/mixins.py
new file mode 100644
index 000000000..877a11039
--- /dev/null
+++ b/numpy/lib/mixins.py
@@ -0,0 +1,167 @@
+"""Mixin classes for custom array types that don't inherit from ndarray."""
+from __future__ import division, absolute_import, print_function
+
+import sys
+
+from numpy.core import umath as um
+
+# None of this module should be exposed in top-level NumPy module.
+__all__ = []
+
+
+def _binary_method(ufunc):
+    def func(self, other):
+        try:
+            if other.__array_ufunc__ is None:
+                return NotImplemented
+        except AttributeError:
+            pass
+        return self.__array_ufunc__(ufunc, '__call__', self, other)
+    return func
+
+
+def _reflected_binary_method(ufunc):
+    def func(self, other):
+        try:
+            if other.__array_ufunc__ is None:
+                return NotImplemented
+        except AttributeError:
+            pass
+        return self.__array_ufunc__(ufunc, '__call__', other, self)
+    return func
+
+
+def _inplace_binary_method(ufunc):
+    def func(self, other):
+        result = self.__array_ufunc__(
+            ufunc, '__call__', self, other, out=(self,))
+        if result is NotImplemented:
+            raise TypeError('unsupported operand types for in-place '
+                            'arithmetic: %s and %s'
+                            % (type(self).__name__, type(other).__name__))
+        return result
+    return func
+
+
+def _numeric_methods(ufunc):
+    return (_binary_method(ufunc),
+            _reflected_binary_method(ufunc),
+            _inplace_binary_method(ufunc))
+
+
+def _unary_method(ufunc):
+    def func(self):
+        return self.__array_ufunc__(ufunc, '__call__', self)
+    return func
+
+
+class NDArrayOperatorsMixin(object):
+    """Mixin defining all operator special methods using __array_ufunc__.
+
+    This class implements the special methods for almost all of Python's
+    builtin operators defined in the `operator` module, including comparisons
+    (``==``, ``>``, etc.) and arithmetic (``+``, ``*``, ``-``, etc.), by
+    deferring to the ``__array_ufunc__`` method, which subclasses must
+    implement.
+
+    This class does not yet implement the special operators corresponding
+    to ``divmod``, unary ``+`` or ``matmul`` (``@``), because these operation
+    do not yet have corresponding NumPy ufuncs.
+
+    It is useful for writing classes that do not inherit from `numpy.ndarray`,
+    but that should support arithmetic and numpy universal functions like
+    arrays as described in :ref:`A Mechanism for Overriding Ufuncs
+    <neps.ufunc-overrides>`.
+
+    As an trivial example, consider this implementation of an ``ArrayLike``
+    class that simply wraps a NumPy array and ensures that the result of any
+    arithmetic operation is also an ``ArrayLike`` object::
+
+        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
+                    # and superclass instances of this type
+                    if not (isinstance(x, self._HANDLED_TYPES) or
+                            isinstance(self, type(x))):
+                        return NotImplemented
+
+                # Defer to the implementation of the ufunc on unwrapped values
+                inputs = tuple(x.value if isinstance(self, type(x)) else x
+                               for x in inputs)
+                if out:
+                    kwargs['out'] = tuple(
+                        x.value if isinstance(self, type(x)) 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)
+
+    In interactions between ``ArrayLike`` objects and numbers or numpy arrays,
+    the result is always another ``ArrayLike``:
+
+        >>> x = ArrayLike([1, 2, 3])
+        >>> x - 1
+        ArrayLike(array([0, 1, 2]))
+        >>> 1 - x
+        ArrayLike(array([ 0, -1, -2]))
+        >>> np.arange(3) - x
+        ArrayLike(array([-1, -1, -1]))
+        >>> x - np.arange(3)
+        ArrayLike(array([1, 1, 1]))
+
+    Note that unlike ``numpy.ndarray``, ``ArrayLike`` does not allow operations
+    with arbitrary, unrecognized types. This ensures that interactions with
+    ArrayLike preserve a well-defined casting hierarchy.
+    """
+
+    # comparisons don't have reflected and in-place versions
+    __lt__ = _binary_method(um.less)
+    __le__ = _binary_method(um.less_equal)
+    __eq__ = _binary_method(um.equal)
+    __ne__ = _binary_method(um.not_equal)
+    __gt__ = _binary_method(um.greater)
+    __ge__ = _binary_method(um.greater_equal)
+
+    # numeric methods
+    __add__, __radd__, __iadd__ = _numeric_methods(um.add)
+    __sub__, __rsub__, __isub__ = _numeric_methods(um.subtract)
+    __mul__, __rmul__, __imul__ = _numeric_methods(um.multiply)
+    if sys.version_info.major < 3:
+        # Python 3 uses only __truediv__ and __floordiv__
+        __div__, __rdiv__, __idiv__ = _numeric_methods(um.divide)
+    __truediv__, __rtruediv__, __itruediv__ = _numeric_methods(um.true_divide)
+    __floordiv__, __rfloordiv__, __ifloordiv__ = _numeric_methods(
+        um.floor_divide)
+    __mod__, __rmod__, __imod__ = _numeric_methods(um.mod)
+    # TODO: handle the optional third argument for __pow__?
+    __pow__, __rpow__, __ipow__ = _numeric_methods(um.power)
+    __lshift__, __rlshift__, __ilshift__ = _numeric_methods(um.left_shift)
+    __rshift__, __rrshift__, __irshift__ = _numeric_methods(um.right_shift)
+    __and__, __rand__, __iand__ = _numeric_methods(um.bitwise_and)
+    __xor__, __rxor__, __ixor__ = _numeric_methods(um.bitwise_xor)
+    __or__, __ror__, __ior__ = _numeric_methods(um.bitwise_or)
+
+    # unary methods
+    __neg__ = _unary_method(um.negative)
+    __abs__ = _unary_method(um.absolute)
+    __invert__ = _unary_method(um.invert)
diff --git a/numpy/lib/tests/test_mixins.py b/numpy/lib/tests/test_mixins.py
new file mode 100644
index 000000000..f45a3c661
--- /dev/null
+++ b/numpy/lib/tests/test_mixins.py
@@ -0,0 +1,189 @@
+from __future__ import division, absolute_import, print_function
+
+import numbers
+import operator
+import sys
+
+import numpy as np
+from numpy.testing import (
+    TestCase, run_module_suite, assert_, assert_equal, assert_raises)
+
+
+PY2 = sys.version_info.major < 3
+
+
+# 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
+            # and superclass instances of this type
+            if not (isinstance(x, self._HANDLED_TYPES) or
+                    isinstance(self, type(x))):
+                return NotImplemented
+
+        # Defer to the implementation of the ufunc on unwrapped values
+        inputs = tuple(x.value if isinstance(self, type(x)) else x
+                       for x in inputs)
+        if out:
+            kwargs['out'] = tuple(
+                x.value if isinstance(self, type(x)) 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 _assert_equal_type_and_value(result, expected, err_msg=None):
+    assert_equal(type(result), type(expected), err_msg=err_msg)
+    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)
+
+
+class TestNDArrayOperatorsMixin(TestCase):
+
+    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):
+            """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_unary_methods(self):
+        array = np.array([-1, 0, 1, 2])
+        array_like = ArrayLike(array)
+        for op in [operator.neg,
+                   # pos is not yet implemented
+                   abs,
+                   operator.invert]:
+            _assert_equal_type_and_value(op(array_like), ArrayLike(op(array)))
+
+    def test_binary_methods(self):
+        array = np.array([-1, 0, 1, 2])
+        array_like = ArrayLike(array)
+        operators = [
+            operator.lt,
+            operator.le,
+            operator.eq,
+            operator.ne,
+            operator.gt,
+            operator.ge,
+            operator.add,
+            operator.sub,
+            operator.mul,
+            operator.truediv,
+            operator.floordiv,
+            # TODO: test div on Python 2, only
+            operator.mod,
+            # divmod is not yet implemented
+            pow,
+            operator.lshift,
+            operator.rshift,
+            operator.and_,
+            operator.xor,
+            operator.or_,
+        ]
+        for op in operators:
+            expected = ArrayLike(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_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):
+        def check(result):
+            assert_(type(result) is tuple)
+            assert_equal(len(result), 2)
+            mantissa, exponent = np.frexp(2 ** -3)
+            _assert_equal_type_and_value(result[0], ArrayLike(mantissa))
+            _assert_equal_type_and_value(result[1], ArrayLike(exponent))
+
+        check(np.frexp(ArrayLike(2 ** -3)))
+        check(np.frexp(ArrayLike(np.array(2 ** -3))))
+
+
+if __name__ == "__main__":
+    run_module_suite()