Call decorator only once in a nested class function

Question

My class has the "public" functions attached with a decorator. In one of the functions, I call another class function that also has the same decorator. But given the decorator was already invoked, I'd like to skip this decorator call. Is there a way to achieve this?

Background: I have some virtual "circuit breakers" that I need to turn turn on/off between the calls. But if I call a function that already touched the circuit breakers once, then I'd want to avoid it in a nested call.

class foo:

    def my_decorator(func):

        @wraps(func)
        def _wrapper(self, *args, **kwargs):
            print("before")
            val = func(self, *args, **kwargs)
            print("after")
            return val
        return _wrapper

    @my_decorator
    def bar(self):
        return 0

    @my_decorator
    def baz(self):
        self.bar()
        return 1

In this example, I see :

f.baz()
before
before
after
after

How do I modify it so I only before and after once, like how bar does:

f.bar()
before
after

Answer 1

My class has the "public" functions attached with a decorator. In one of the functions, I call another class function that also has the same decorator. But given the decorator was already invoked, I'd like to skip this decorator call. Is there a way to achieve this?

So - first - keeping simple things simple: if this is a one-off situation, and you can hard-code inside baz that you want to call bar without the decorator, you can simply explicitly call the function given in its __wrapped__ attribut, like in Karl's answer. The only objective drawback is that it won't work in conjunction with other decorators.

However, this approach involve knowing, and remembering, at development time, which other methods are decorated with the decorator we want to run only once, and write the inner code in accord to that. When working with decorators, the more standard practice is that once you add your decorator as a modifier to your function or method, you no longer have to worry about it - including not having to modify your code to cooperate with the decorator.

And then, is there a way for, inside the decorator code modify the decorated function so that the decorator will be skipped when another decorated method will be called?

No. Not an feasible and sane way, nonetheless - to avoid the decorator to be entered at all could possibly only be done by modifying the bytecode in the decorated function, so that it could test when a call is to a sibling-decorated function and call the wrapped function instead.

But once the decorator is entered, testing for it becomes a bit easier - you then just refrain from re-running the code you don't want to run twice (or more times) - this is not trivial, but it is for certain feasible in a reliable way.

For one thing, your decorator needs to share a state, so that it knows when it is "re-entering" and skip relevant code. One good way to share this might be through Python's context-vars: these are states that are unique to a thread or a task-callstack in async based code.

If however, your code might involve calling methods in other instances of the same object, and the decorator code needs to run once in each instance, you might need some more state keeping. I won't deal with this case in the proof of concept bellow.

So a first, verbose but safe, approach is this:

import contextvars

from functools import wraps

deco_state = contextvars.ContextVar("deco_state")

def my_decorator(func):

    @wraps(func)
    def _wrapper(self, *args, **kwargs):
        try:
            deco_state.set(depth:=(deco_state.get(0) + 1))
            if depth <= 1:
                print("before")
            val = func(self, *args, **kwargs)
            if depth <= 1:
                print("after")
        finally:
            deco_state.set(deco_state.get() - 1)
        return val
    return _wrapper

class Foo:

    @my_decorator
    def bar(self):
        return 0

    @my_decorator
    def baz(self):
        self.bar()
        return 1

Foo().baz()

As stated: this will be resilient to multi-threaded and multi-async-task - but only if the different tasks are calling methods on the same instance. given this design, I suppose you want the decorator code to run once, and only once, in each instance of your Foo class, regardless of the method been called from different threads or different tasks in async code.

Then, instead of context-vars, one might just key the call counter to the instance.

Also, the example above is cumbersome, requiring two if statements and modifying the decorator code itself. We can take the need for rework this, and create a class that can decorate your wrapper (a decorator for the decorator), and be called in place of the functools.wraps decorator: it then can take the decision of running your wrapper or calling the decorated function directly.

So this class also performs the role of the outer-decorator code, the one which just annotates the func in a closure for the "wrapper" function, which is the real decorator.

Upon being called, the wrapper provided by this class will then know when to run the original function through your decorator, or go directly to the original function, and use contextvars.Context to check for recursivity. Contextvars are a relatively new thing in Python, and they can track the "call stack" i.e. if one thread called "baz" in one context, another thread can call the same function in other contexts, and code inside those methods (or the decorator method), can see distinct values in the ContextVar instance. In this case we just need a boolean to know if the decorator has run in a given call-stack.

I took some time to make this production quality, as this code can also solve a problem I've met more than once: how to avoid a decorator to run twice, as you want, but when one overrides a decorated method in a subclass.

With this code, all one needs to do is to decorate both the method on the superclass and the overriden method in the child-classes with the same decorator.

All in all, this is the total code, with your example modified to run at the end:

from functools import wraps
from threading import RLock
import contextvars

RUNONCE = contextvars.ContextVar("RUNONCE", default=None)

class RunOnceDecorator:

    def __init__(self, decorator):
        self.decorator = decorator

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            v = RUNONCE.get()
            try:
                if v is None:
                    def mid():
                        # intermediary function needed to set a state on contextvars.context
                        nonlocal result
                        RUNONCE.set(True)
                        result = self.decorator(func,*args, **kwargs)
                    ctx = contextvars.Context()
                    ctx.run(mid)
                else:
                    ctx = None
                    result = func(*args, **kwargs)
            finally:
                del ctx
            return result
        return wrapper


# so, this is the thing: we can supress the outermost
# function, and get `func` as the very first
# parameter whenever it is called:

@RunOnceDecorator
def my_decorator(func, self, *args, **kwargs):
    print("before")
    val = func(self, *args, **kwargs)
    print("after")
    return val

class Foo:

    @my_decorator
    def bar(self):
        return 0

    @my_decorator
    def baz(self):
        self.bar()
        return 1

Foo().baz()

This will output simply:

$ python blip2.py 
before
after

There are no tests here, but it should work against multiple instances and multi-threading.