pythonpython-decoratorsinspect

A reliable way to check if a method has been wrapped with a decorator from within it


In python 3.7 or higher.

I have the following class:

def dec(method):
    @functools.wraps(method)
    def wrapper(*args, **kwargs):
        print("wrapped")
        return method(*args, **kwargs)
    return wrapper

class A:
    @dec()
    def f1(self):
        print(is_wrapped())

    def f2(self):
        print(is_wrapped())

I want A().f1() to print True and A().f2() to print False.

I created the following code for is_wrapped:

def is_wrapped():
    frame = inspect.currentframe().f_back
    for v in frame.f_back.f_locals.values():
        if hasattr(v, '__code__') and v.__code__ is frame.f_code:
            return True

    return False

While this seems to work it can fail if the caller of f2 has a local variable that contains f2 but does not decorate it.

For my specific case you can assume the following.

  1. The is_wrapped function is called directly from the function in a way that inspect.currentframe().f_back is the method that called it (i.e, f1 or f2)

  2. The is_wrapped function cannot get any arguments.

  3. the decorator can be implemented in any way as long as it uses functool.wraps for the wrapped function.

Is there any more reliable way to achieve this?


Solution

  • I have revisited this question since my last comment, and here's a probably better way: instead of looking for a local variable with same __code__ as frame's code, let's look for something that has __wrapped__ satisfying that condition.

    Why is this better? Since you say that functools.wraps is a requirement for decorator, we can be sure that a decorator must have some local with __wrapped__ - that's the function built inside. So now we're looking for the following: "traversing from current frame, can the grandparent be a function definition that's called now?"

    So, if current frame is the one inside is_wrapped, we traverse one back to the original method (f1). This assumes that decorator does actually call the decorated function, which makes sense given functools.wraps requirement. If there are no previous frame, we're at some wrong place, probably is_wrapped shouldn't even be called from that. Otherwise, go a step further: we're now in function built inside the decorator.

    So the sketch of a modified version may look like this:

    import gc
    
    def is_wrapped():
        frame = inspect.currentframe().f_back
        if frame is None:
            raise RuntimeError("No call frame found")
        orig_code = frame.f_code  # Looking for this original definition
        if frame.f_back is None:
            return False
        deco_code = frame.f_back.f_code
        
        for o in gc.get_referrers(frame.f_back.f_code):
            if getattr(o, '__code__', None) is deco_code and hasattr(o, '__wrapped__'):
                return getattr(o.__wrapped__, '__code__', None) is orig_code 
                
        return False
    

    It passes all "testcases" from your question and comments, works for "deep" decorators that accept parameters, but is still breakable - for example, if a decorator-created function does not call the original code, you're out of luck. But it's something less trivial than simply referencing the method in enclosing scope, isn't it?

    Credits for gc approach of retrieving the actual function go to @MikeHordecki.

    Use at your own risk and test thoroughly. Also note that

    Warning: Care must be taken when using objects returned by get_referrers() because some of them could still be under construction and hence in a temporarily invalid state. Avoid using get_referrers() for any purpose other than debugging.