Class QuantityND

lib/matplotlib/tests/test_image.py:1294–1351

# There is just enough implemented to test handling of ndarray
# subclasses.
class QuantityND(np.ndarray):
    def __new__(cls, input_array, units):
        obj = np.asarray(input_array).view(cls)
        obj.units = units
        return obj

    def __array_finalize__(self, obj):
        self.units = getattr(obj, "units", None)

    def __getitem__(self, item):
        units = getattr(self, "units", None)
        ret = super().__getitem__(item)
        if isinstance(ret, QuantityND) or units is not None:
            ret = QuantityND(ret, units)
        return ret

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        func = getattr(ufunc, method)
        if "out" in kwargs:
            return NotImplemented
        if len(inputs) == 1:
            i0 = inputs[0]
            unit = getattr(i0, "units", "dimensionless")
            out_arr = func(np.asarray(i0), **kwargs)
        elif len(inputs) == 2:
            i0 = inputs[0]
            i1 = inputs[1]
            u0 = getattr(i0, "units", "dimensionless")
            u1 = getattr(i1, "units", "dimensionless")
            u0 = u1 if u0 is None else u0
            u1 = u0 if u1 is None else u1
            if ufunc in [np.add, np.subtract]:
                if u0 != u1:
                    raise ValueError
                unit = u0
            elif ufunc == np.multiply:
                unit = f"{u0}*{u1}"
            elif ufunc == np.divide:
                unit = f"{u0}/({u1})"
            elif ufunc in (np.greater, np.greater_equal,
                           np.equal, np.not_equal,
                           np.less, np.less_equal):
                # Comparisons produce unitless booleans for output
                unit = None
            else:
                return NotImplemented
            out_arr = func(i0.view(np.ndarray), i1.view(np.ndarray), **kwargs)
        else:
            return NotImplemented
        if unit is None:
            out_arr = np.array(out_arr)
        else:
            out_arr = QuantityND(out_arr, unit)
        return out_arr

    @property
    def v(self):
        return self.view(np.ndarray)