A wrapper for dictionaries with numeric values

NumericDict.py

_numeric_types = (int, float, complex)


class NumericDict(dict):
    def __add__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: v + other)
        if isinstance(other, NumericDict):
            return self.__class__({k: v + other.get(k, 0) for k, v in self.items()})

        raise NotImplementedError()

    def __radd__(self, other):
        return self + other

    def __sub__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: v - other)
        if isinstance(other, NumericDict):
            return self.__class__({k: v - other.get(k, 0) for k, v in self.items()})

        raise NotImplementedError()

    def __rsub__(self, other):
        return -self + other

    def __mul__(self, other):
        if isinstance(other, (int, float)):
            return self.apply(lambda v: v * other)
        if isinstance(other, NumericDict):
            return self.__class__({k: v * other.get(k, 1) for k, v in self.items()})

        raise NotImplementedError()

    def __rmul__(self, other):
        return self * other

    def __truediv__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: v / other)
        if isinstance(other, NumericDict):
            return self.__class__({k: v / other.get(k, 1) for k, v in self.items()})

        raise NotImplementedError()

    def __rtruediv__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: other / v)

        raise NotImplementedError()

    def __floordiv__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: v // other)
        if isinstance(other, NumericDict):
            return self.__class__({k: v // other.get(k, 1) for k, v in self.items()})

        raise NotImplementedError()

    def __rfloordiv__(self, other):
        if isinstance(other, _numeric_types):
            return self.apply(lambda v: other // v)

        raise NotImplementedError()

    def __neg__(self):
        return self.__class__({k: -v for k, v in self.items()})

    def apply(self, func):
        return self.__class__({k: func(v) for k, v in self.items()})