Search code examples
pythonpython-3.xpython-decorators

How to add attributes to decorated functions when combining decorators?


Two function decorators are defined. The target is to detect if 0, 1, or 2 decorators are applied to a function.

Why does the below code return "False" for the 2nd decorator?

def decorator1(f):
    def wrapped(*args, **kwargs):
        f(*args, **kwargs)

    wrapped.dec1 = True
    return wrapped


def decorator2(f):
    def wrapped(*args, **kwargs):
        f(*args, **kwargs)

    wrapped.dec2 = True
    return wrapped


@decorator1
@decorator2
def myfunc():
    print(f"running myfunc")


if __name__ == "__main__":
    myfunc()

    print(f"myfunc has decorator1: {getattr(myfunc, 'dec1', False)}")
    print(f"myfunc has decorator2: {getattr(myfunc, 'dec2', False)}")

Result:

running myfunc
myfunc has decorator1: True 
myfunc has decorator2: False

I am using Python 3.9.


Solution

  • Actually this sentence in print statement is not correct: "myfunc has...". No. myfunc has neither dec1 nor dec2. What you get back from calling your decorator is a new function not the function you passed. I'll explain it.

    This code:

    @decorator1
    @decorator2
    def myfunc():
        print(f"running myfunc")
    

    is equivalent to:

    def myfunc():
        print('running myfunc')
    
    myfunc = decorator2(myfunc)
    myfunc = decorator1(myfunc)
    

    Remember that you add those dynamic attributes to the wrapped functions. They don't both apply to the same myfunc function object. They apply to "their" wrapped function. The wrapped functions are two different objects in those two decorators:

    myfunc = decorator2(myfunc)
    print(myfunc)   # <function decorator2.<locals>.wrapped at 0x1028456c0>
    myfunc = decorator1(myfunc)
    print(myfunc)   # <function decorator1.<locals>.wrapped at 0x102845760>
    

    After stacking decorators, you can expect the decorator1.<locals>.wrapped to have dec1(because it's the outer one) but it doesn't have dec2. It's on the decorator2.<locals>.wrapped:

    def decorator1(f):
        def wrapped(*args, **kwargs):
            return f(*args, **kwargs)
        wrapped.dec1 = True
        return wrapped
    
    
    def decorator2(f):
        def wrapped(*args, **kwargs):
            return f(*args, **kwargs)
        wrapped.dec2 = True
        return wrapped
    
    
    def myfunc():
        print('running myfunc')
    
    
    myfunc = decorator2(myfunc)
    myfunc = decorator1(myfunc)
    
    
    print(f"{getattr(myfunc, 'dec1', False)}")  # True
    print(f"{getattr(myfunc.__closure__[0].cell_contents, 'dec2', False)}")  # True
    

    If you had these kinds of decorators, your assumptions would work:

    def decorator1(f):
        f.dec1 = True
        return f
    
    def decorator2(f):
        f.dec2 = True
        return f
    
    @decorator1
    @decorator2
    def myfunc():
        print('running myfunc')
    
    print(f"{getattr(myfunc, 'dec1', False)}")  # True
    print(f"{getattr(myfunc, 'dec2', False)}")  # True
    

    How can you achieve it with your existing code?

    I don't think it's a clean way but it works if you want:

    def decorator1(f):
        def wrapped(*args, **kwargs):
            return f(*args, **kwargs)
    
        while hasattr(f, '__wrapped__'):
            f = f.__wrapped__
        f.dec1 = True
        wrapped.__wrapped__ = f
        return wrapped
    
    def decorator2(f):
        def wrapped(*args, **kwargs):
            return f(*args, **kwargs)
    
        while hasattr(f, '__wrapped__'):
            f = f.__wrapped__
        f.dec2 = True
        wrapped.__wrapped__ = f
        return wrapped
    
    @decorator1
    @decorator2
    def myfunc():
        print('running myfunc')
    
    print(f"{getattr(myfunc.__wrapped__, 'dec1', False)}")  # True
    print(f"{getattr(myfunc.__wrapped__, 'dec2', False)}")  # True
    

    You see, with that while loop, I go deep inside to find the original myfunc and add it to the returned wrapped function. Also I add the decX attributes to the original myfunc.