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 defe
| 57 | |
| 58 | |
| 59 | class NDArrayOperatorsMixin: |
| 60 | """Mixin defining all operator special methods using __array_ufunc__. |
| 61 | |
| 62 | This class implements the special methods for almost all of Python's |
| 63 | builtin operators defined in the `operator` module, including comparisons |
| 64 | (``==``, ``>``, etc.) and arithmetic (``+``, ``*``, ``-``, etc.), by |
| 65 | deferring to the ``__array_ufunc__`` method, which subclasses must |
| 66 | implement. |
| 67 | |
| 68 | It is useful for writing classes that do not inherit from `numpy.ndarray`, |
| 69 | but that should support arithmetic and numpy universal functions like |
| 70 | arrays as described in `A Mechanism for Overriding Ufuncs |
| 71 | <https://numpy.org/neps/nep-0013-ufunc-overrides.html>`_. |
| 72 | |
| 73 | As an trivial example, consider this implementation of an ``ArrayLike`` |
| 74 | class that simply wraps a NumPy array and ensures that the result of any |
| 75 | arithmetic operation is also an ``ArrayLike`` object:: |
| 76 | |
| 77 | class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin): |
| 78 | def __init__(self, value): |
| 79 | self.value = np.asarray(value) |
| 80 | |
| 81 | # One might also consider adding the built-in list type to this |
| 82 | # list, to support operations like np.add(array_like, list) |
| 83 | _HANDLED_TYPES = (np.ndarray, numbers.Number) |
| 84 | |
| 85 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): |
| 86 | out = kwargs.get('out', ()) |
| 87 | for x in inputs + out: |
| 88 | # Only support operations with instances of _HANDLED_TYPES. |
| 89 | # Use ArrayLike instead of type(self) for isinstance to |
| 90 | # allow subclasses that don't override __array_ufunc__ to |
| 91 | # handle ArrayLike objects. |
| 92 | if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)): |
| 93 | return NotImplemented |
| 94 | |
| 95 | # Defer to the implementation of the ufunc on unwrapped values. |
| 96 | inputs = tuple(x.value if isinstance(x, ArrayLike) else x |
| 97 | for x in inputs) |
| 98 | if out: |
| 99 | kwargs['out'] = tuple( |
| 100 | x.value if isinstance(x, ArrayLike) else x |
| 101 | for x in out) |
| 102 | result = getattr(ufunc, method)(*inputs, **kwargs) |
| 103 | |
| 104 | if type(result) is tuple: |
| 105 | # multiple return values |
| 106 | return tuple(type(self)(x) for x in result) |
| 107 | elif method == 'at': |
| 108 | # no return value |
| 109 | return None |
| 110 | else: |
| 111 | # one return value |
| 112 | return type(self)(result) |
| 113 | |
| 114 | def __repr__(self): |
| 115 | return '%s(%r)' % (type(self).__name__, self.value) |
| 116 |