python - Decorating top function in MRO

Question

How can I decorate the last function in a class inheritance?

If I decorate a superclass function, the subclass function overrides the decorator. I'd like to find out if there is a neat way to automatically decorate the top function in the MRO.

def wrapper(f):
    def _wrap(*args, **kwargs):
        print("In wrapper")
        return f(*args, **kwargs)
    return _wrap


class A:

    @wrapper
    def f(self):
        print("In class A")


class B(A):
    def f(self):
        print("In class B")


if __name__ == '__main__':
    a = A()
    b = B()
    print("Calling A:")
    a.f()
    print("Calling B:")
    b.f()

Here is the output. As expected, B.f() does not call the wrapper, though I'd like it to.

Calling A:
In wrapper
In class A

Calling B:
In class B

Here is what I have tried thus far. A metaclass that holds all the decorators and injects them during class instantiation.

from abc import ABCMeta


class WrapperMetaClass(ABCMeta):
    def __init__(cls, *args, **kwargs):
        wrappers_dict = getattr(cls, "_wrappers")
        for attr_name in dir(cls):
            if attr_name not in wrappers_dict:
                continue
            else:
                wrapper = wrappers_dict[attr_name]
            attr = getattr(cls, attr_name)
            if not hasattr(attr, '__call__'):
                raise Exception("What you're trying to wrap is not a function!")
            attr = wrapper(attr)
            setattr(cls, attr_name, attr)

        super().__init__(*args, **kwargs)

This works:

class A(metaclass=WrapperMetaClass):
    _wrappers = {
        "f": wrapper
    }

    def f(self):
        print("In class A")


class B(A):
    def f(self):
        print("In class B")

The output is what I wanted.

Calling A:
In wrapper
In class A

Calling B:
In wrapper
In class B

However, this runs into a different issue. If B does not override f, the metaclass wraps A.f() twice. This makes sense, as both A and B inherit WrapperMetaClass, so A.f() is wrapped first, and then B.f() is wrapped again.

class A(metaclass=WrapperMetaClass):
    _wrappers = {
        "f": wrapper
    }

    def f(self):
        print("In class A")


class B(A):
    pass

The output becomes:

Calling A:
In wrapper
In class A

Calling B:
In wrapper
In wrapper
In class A

I have no idea what else I could do.

question from:https://stackoverflow.com/questions/65950552/decorating-top-function-in-mro

Answer 1

Yes, I remember facing this once or twice - and you are on the right track.

But first things first: if the logic in your "wrapper" is something that could be put in a method in the base class, then breaking-up the methods in smaller-tasks, and have a "method slot" system is preferable to this, as user 2357112 supports monica puts in the comments. If you find out you really need or prefer decorators, the full code is bellow

class A:
    def do_things(self):
        create_connection()  # <- this is the code you'd are putting in the wrapper in the other approach
        do_thing_1()

class B(A):
    def do_things(self):
        # here we have to do thing_1 and thing_2, but
        # the connection is created in the superclass method... 
        # this could be the origin of your question
        
# Refactor to:

class A:
    def do_things(self):
        create_connection()
        self.internal_do_things()
        
    def internal_do_things(self):
        do_thing_1()
        
class B(A):
    def internal_do_things(self):
        super().internal_do_things()
        do_thing_2()

So, classical inheritance and OO solves this

---

If you need the decorators anway:

The thing to do is to have the decorator itself, the "wrapper", get a way to "know" if it already was called in an outer method (i.e. a method in a subclass which calls super()), and just act as a transparent wrapper in this case.

It gets a bit further complicated when we want a robust solution: a wrapper that can work for different methods in the same class, and does not get confused if they are called concurrently (in different threads, or a method calling another method, not super(), which should trigger the wrapper).

And in the end, the mechanisms for that are complicated enough that they should not get in the way of your actual wrapper - so, ideally they should be built as a decorator themselves, which will decorate your wrapper.

[hours later] So, sorry if it does not look "neat" - it turns out implementing what is described above is a bit more complex than I thought initially - we need an intermediate decorator level (called meta_wrapper_applier in the code), so that the metaclass can re-wrap the methods each time they are redeclared.

I hope the comments in the code and variable names are enough to understand the idea:

from abc import ABCMeta
from functools import wraps
import threading


class WrapperMetaClass(ABCMeta):
    def __init__(cls, name, bases, ns, **kw):
        super().__init__(name, bases, ns, **kw)

        # Get the wrapped methods for all the superclasses
        super_wrappers = {}
        for supercls in cls.__mro__[::-1]:
            super_wrappers.update(supercls.__dict__.get("_wrappers", {}))

        # unconditionally install a wrappers dict for each subclass:
        sub_wrappers =  cls._wrappers = {}

        for attrname, attr in ns.items():
            if attrname in super_wrappers:
                # Applies the wrapper in the baseclass to the subclass method:
                setattr(cls, attrname, super_wrappers[attrname]._run_once_wrapper(attr))
            elif hasattr(attr, "_run_once_wrapper"):
                # Store the wrapper information in the cls for use of the subclasses:
                sub_wrappers[attrname] = attr


def run_once_method_decorator(original_wrapper):
    re_entering_stacks = {}

    # This is the callable used to place a wrapper on the original
    # method and on each overriden method.
    # All methods with the same name in the subclasses  will share the same original wrapper and the
    # "re_entering_stacks" data structure.
    def meta_wrapper_applier(raw_method):
        wrapped_method_in_subclass = None
        @wraps(original_wrapper)
        def meta_wrapper(*args, **kw):
            nonlocal wrapped_method_in_subclass

            # uses a plain list to keep track of re-entering the same-named method
            # in each thread:
            re_entering_stack = re_entering_stacks.setdefault(threading.current_thread(), [])
            re_entering = bool(re_entering_stack)
            try:
                re_entering_stack.append(1)
                if re_entering:
                    result = raw_method(*args, **kw)
                else:
                    if wrapped_method_in_subclass is None:
                        # Applies the original decorator lazily, and caches the result:
                        wrapped_method_in_subclass = original_wrapper(raw_method)

                    result = wrapped_method_in_subclass(*args, **kw)
            finally:
                re_entering_stack.pop()
            return result
        # registry = original_wrapper.__dict__.setdefault("_run_once_registry", {})


        meta_wrapper._run_once_wrapper = meta_wrapper_applier
        return meta_wrapper

    return meta_wrapper_applier

# From here on, example code only;

@run_once_method_decorator
def wrapper(f):
    @wraps(f)
    def _wrap(*args, **kwargs):
        print("Entering wrapper")
        result = f(*args, **kwargs)
        print("Leaving wrapper
")
        return result
    return _wrap

@run_once_method_decorator
def other_wrapper(f):
    @wraps(f)
    def _wrap(*args, **kwargs):
        print("Entering other wrapper")
        result = f(*args, **kwargs)
        print("Leaving other wrapper
")
        return result
    return _wrap


class A(metaclass=WrapperMetaClass):

    @wrapper
    def f(self):
        print("In class A")

    def g(self):
        print("g in A")

class B(A):
    def f(self):
        print("In class B")
        super().f()

    @other_wrapper
    def g(self):
        print("g in B")
        super().g()

class C(B):
    def g(self):
        print("g in C")
        super().g()


if __name__ == '__main__':
    a = A()
    b = B()
    print("Calling A:")
    a.f()
    a.g()
    print("Calling B:")
    b.f()
    b.g()
    print("Calling C:")
    C().g()

Output:

Calling A:
Entering wrapper
In class A
Leaving wrapper

g in A
Calling B:
Entering wrapper
In class B
In class A
Leaving wrapper

Entering other wrapper
g in B
g in A
Leaving other wrapper

Calling C:
Entering other wrapper
g in C
g in B
g in A
Leaving other wrapper