Re-type operator overloads when subclassing a builtin or stdlib type

When subclassing native types such as float or dict, or stdlib types such as Counter, it’s often surprising or frustrating for new users that Subclass() + Subclass() is an instance of the native type rather than the subclass. (It is so for performance reasons I assume, so I’m not proposing we change that.)

Mitigating that behavior and making the new class return instances of itself can be written in several ways. Here (source) is my personal favorite, limiting code duplication :

def absolute_wrap(func):
    """
    Wraps func_name into a method of absolute. The wrapped method
    converts a float result back to absolute.
    """

    def wrapper(*args):
        rv = func(*args)

        if rv is NotImplemented:
            return rv
        else:
            return absolute(rv)

    return wrapper

class absolute(float):
    """
    This represents an absolute float coordinate.
    """

    slots = ()

    def __repr__(self):
        return "absolute({})".format(float.__repr__(self))

    def __divmod__(self, value):
        return self//value, self%value

    def __rdivmod__(self, value):
        return value//self, value%self

for fn in (
    '__abs__',
    '__add__',
    # '__bool__', # non-float
    '__ceil__',
    # '__divmod__', # special-cased above, tuple of floats
    # '__eq__', # non-float
    '__floordiv__',
    # '__format__', # non-float
    # '__ge__', # non-float
    # '__gt__', # non-float
    # '__hash__', # non-float
    # '__int__', # non-float
    # '__le__', # non-float
    # '__lt__', # non-float
    '__mod__',
    '__mul__',
    # '__ne__', # non-float
    '__neg__',
    '__pos__',
    '__pow__',
    '__radd__',
    # '__rdivmod__', # special-cased above, tuple of floats
    '__rfloordiv__',
    '__rmod__',
    '__rmul__',
    '__round__',
    '__rpow__',
    '__rsub__',
    '__rtruediv__',
    # '__str__', # non-float
    '__sub__',
    '__truediv__',
    # '__trunc__', # non-float

    # 'as_integer_ratio', # tuple of non-floats
    'conjugate',
    'fromhex',
    # 'hex', # non-float
    # 'is_integer', # non-float
):
    f = getattr(float, fn)
    setattr(absolute, fn, absolute_wrap(f))

del absolute_wrap, fn, f

(EDIT : I had forgotten to special-case __rdivmod__ as well - which shows you how much we would benefit from an integrated solution)

As you can see, it’s still quite tedious to do.
I don’t have a clear plan of how it could be made better, but I think maybe a class decorator similar to functools.total_ordering which would do that automatically from the newly created class ? I thought of a super-like assignment in the class body, method by method, but it would possibly be harder to implement and would still require every method name to be written twice. I’m open to alternative solutions.

The question of which methods the decorator should wrap is a bit trickier : should it wrap all datamodel dunders (including getitem for list and tuple subclasses, which should return instances of the new class when passed slices) with some being blacklisted by default as I did in the code above ?

The way I picture it would be this way : type_wrap(*, include=(), exclude=()). Noting datamo_meth the set of methods I listed above, it would wrap datamo_meth.difference(cls.__dict__, exclude).union(include), and then special-case divmod and getitem as written above. That way, the above code could be simplified as :

@functools.type_wrap(include=("conjugate", "fromhex"))
class absolute(float):
    __slots__ = ()
    def __repr__(self):
        return "absolute({})".format(float.__repr__(self))

If I want, say, __add__ to work in a specific way, I write it in the class’s body and because of the cls.__dict__ exclusion above, the decorator would not overwrite it. If I want __sub__ to still return a float (for some reason), I can pass exclude=("__sub__",) to the decorator.

I imagine this would go to benchmark implementation and several other steps before being merged in the stdlib / in the functools module, but I’m looking for API feedback, obvious footguns I would have missed, alternative solutions, et cetera.

4 Likes

Here is an example implementation. I’m still not sure what to do with __round__, but the rest seems to work very well.
I figured out what to do with __round__. The code seems ready for prod, as far as I tested. Review and comments would be appreciated.

This has been made an issue and a related pull request.