Search code examples
pythonpython-3.xpython-decorators

Getting rid of explicit super


I'd like to implement something like this

def after(f):
    def inner(*args, **kwargs):
        super().f(*args, **kwargs)
        f(*args, **kwargs)
    return inner


class A:
  def f(self):
    print ('hello')


class B(A):

  @after
  def f(self):
    print ('world')

b = B()
b.f()

that is I would like to get rid of explicit super in some of my classes and replace it with @before / @after decorator (maybe with parameters).

That is, in this example, I would like hello world to be printed.

the idea is to increase the readability of the code, as in some classes I often use multiple inheritance, so I often override methods and often have to use super().

I think I could use inspect to determine the class instance that calls the decorator (although not sure about performance if I have many class instances).

is there a way to do this without sacrificing performance?


Solution

  • You can make your decorator work, you just need it to make it a descriptor class, rather than a function. You need to implement the __set_name__ method to get a reference to the class you've been added to. With the class reference, you can make a two-argument super call:

    import functools
    
    class after:
        def __init__(self, method):
            self.method = method
    
        def __set_name__(self, owner, name):
            self.owner = owner
            self.name = name                 # using self.method.__name__ might be better?
    
        def __get__(self, instance, owner):
            if instance is None:
                return self
            return functools.partial(self, instance)
    
        def __call__(self, instance, *args, **kwargs):
            assert(self.owner is not None and self.name is not None)
            getattr(super(self.owner, instance), self.name)(*args, **kwargs)
            return self.method(instance, *args, **kwargs)
    

    You could do a before too, which would be nearly the same, just with the last two lines in the reverse order (and some fiddling to handle the return value).

    I'd note that this decorator is quite a bit less generally useful than calling super the normal way since you can't usefully interact with the value returned by the overridden method, or change the arguments being passed in to it. There's no before or after decorated method that can replicate these classes:

    class Foo:
        def foo(self, x, y):
            return x + y
    
    class Bar(Foo):
        def foo(self, x, y, z):
            return super().foo(x//2, y+1) * z