Poshi
Poshi

Reputation: 5762

Python decorators passing attributes

I'm using decorators to enhance some methods, but I'm suffering from the lack of interoperability between them.

As an example, let's say I want to use the functools.cache decorator to memoize the results and a hand-made decorator to count the number of calls to that method:

from functools import cache, wraps
from typing import Callable


def counted(func: Callable) -> Callable:

    @wraps(func)
    def wrapped(*args, **kwargs):
        setattr(wrapped, "calls", getattr(wrapped, "calls") + 1)
        return func(*args, **kwargs)

    setattr(wrapped, "calls", 0)

    return wrapped


@counted
@cache
def func_a(data):
    return data


if __name__ == "__main__":
    func_a(1)
    func_a.clear_cache()
    print(func_a.calls)

The code, as it is shown, fails at the func_a.clear_cache(), because the counted decorator did not pass the methods/attributes that cache added to the function. If we swap the two decorators, then the print(func_a.calls) will fail because the cache decorator did not pass the attribute calls that was set by the inner decorator.

Is there a pythonic way to get a final function that contains every bit that has been added by the decorators?

I know I can modify the counted decorator to explicitly pass the cache-added attributes, but the issue comes when you are using two or more third-party decorators.

Upvotes: 2

Views: 570

Answers (1)

Daniil Fajnberg
Daniil Fajnberg

Reputation: 18458

Decoration is just syntactic sugar for passing some callable (function or class) to another callable (the decorator) and that syntax is restricted to class/function definition statements.

Given some decorator dec, writing

@dec
def f(): ...

is equivalent to this:

def f(): ...

f = dec(f)

It is also important to stress that nothing inherently special is happening by virtue of decoration. The following is totally valid (albeit not very useful):

def dec(_): return 1

class Foo: pass

@dec
class Bar: pass

def f(): pass

@dec
def g(): pass

print(Foo)  # <class '__main__.Foo'>
print(Bar)  # 1
print(f)    # <function f at 0x7fdf...>
print(g)    # 1

This goes to show that there is nothing magical about decoration leaving some sort of "trace" on the output of the decorator.

The Bar class and the g function are essentially consumed by the dec function and since no reference to them is returned by it, they are no longer in any way available after this decoration.

There is also nothing inherently special about returning functions from a decorator:

def f():
    return "There is no spoon"

def dec(_func):
    return f

@dec
def g():
    return "Hi mom"

print(g.__name__)  # f
print(g())         # There is no spoon

Again, the decorator is just a function and in this case it returns another function, but nothing in this process does anything magical (or anything whatsoever) with the function g. In this example it is basically lost after decoration.

To get to an example more representative of real-world scenarios, decorators are usually written such that they do preserve something about the callable being decorated, but this typically just means that a wrapper function is defined inside the decorator and inside that wrapper the original callable is called.

def dec(func):
    def wrapper():
        return func() + " There is no spoon."
    return wrapper

@dec
def f():
    return "Hi mom."

print(f.__name__)  # wrapper
print(f())         # Hi mom. There is no spoon.

The reference to the original function f is not lost, but it is inside the local namespace of the wrapper returned by dec and there is no way of getting to it anymore.

All of this is to drive the point home why there is no magical built-in way to somehow "preserve" any attributes of the callable being decorated. You need to take care of this yourself, if you want your decorator to do that. The same way you would have to write that kind of logic for any other function that takes some object as its argument, if you expect some attribute of that object to be present in the output of that function. And if you are using someone else's function and they don't do that, you are out of luck.

functools.wraps addresses this by giving us a quasi-standard pattern for writing decorator-wrappers, that keeps an explicit reference to the object being decorated in the __wrapped__ attribute of the wrapper. But nothing forces you to use that pattern and if someone doesn't, again, you are out of luck.

The best thing thing you could do is to write (yet another) custom decorator that relies on other decorators using functools.wraps (or functools.update_wrapper) to recursively propagate everything from the chain of wrapped objects to the top-wrapper. It could look something like this:

from functools import wraps

def propagate_all_attributes(func):
    wrapped = getattr(func, "__wrapped__", None)
    if wrapped is not None:
        propagate_all_attributes(wrapped)
        # Add attributes from `wrapped` that are *not* present in `func`:
        for attr_name, attr_value in wrapped.__dict__.items():
            if attr_name not in func.__dict__:
                func.__dict__[attr_name] = attr_value
    return func

def dec1(func):
    @wraps(func)
    def wrapper():
        return func() + " There is no spoon."
    wrapper.x = 1
    wrapper.y = 2
    return wrapper

def dec2(func):
    @wraps(func)
    def wrapper():
        return func() + " Have a cookie."
    wrapper.y = 42
    return wrapper

@propagate_all_attributes
@dec2
@dec1
def f():
    """Some function"""
    return "Hi mom."

print(f.__name__)  # f
print(f.__doc__)   # Some function
print(f.x)         # 1
print(f.y)         # 42
print(f())         # Hi mom. There is no spoon. Have a cookie.

But again, this will not work, if one of the decorators below it does not (properly) set the __wrapped__ attribute on the object it returns.

That approach would of course allow for additional customization, like e.g. telling your decorator, which attributes to "pull up" from the wrapped object or which to exclude or whether to overwrite attributes set by later decorators with those of inner objects etc..

Assuming you are always able to check the source of third-party decorators you use, you could at least get some of what you are looking for this way, by applying it to decorators that correctly utilize the @wraps-pattern.

Upvotes: 1

Related Questions