diff options
Diffstat (limited to 'numpy/lib/mixins.py')
| -rw-r--r-- | numpy/lib/mixins.py | 167 |
1 files changed, 167 insertions, 0 deletions
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) |