profile

Rodrigo GirΓ£o SerrΓ£o

πŸπŸ“ what you should ALWAYS do when writing decorators

published3 months ago
3 min read

Hey there, how are you doing?

In this issue of the Mathspp Insider we will talk about the one thing you should do every time you define a Python decorator.


functools.wraps

TL;DR: whenever you define a decorator in Python, use functools.wraps to preserve things like the documentation string and the name of the function!

Here is an example of a decorator that prints the arguments passed in to the decorated function:

def print_args(f):
    def wrapper(*args, **kwargs):
        print(f"{args = } and {kwargs = }")
        return f(*args, **kwargs)
    return wrapper

You can apply this decorator to any function you'd like.
Then, that function will print its arguments every time you call it:

>>> @print_args
... def add(a, b):
...     """Adds two numbers together."""
...     return a + b
...
>>> add(3, 4)
args = (3, 4) and kwargs = {}
7
>>> add(-2, b=2)
args = (-2,) and kwargs = {'b': 2}
0

The issue

However, if you inspect the function, it looks all funky:

>>> add
<function print_args..wrapper at 0x000001C58662B310>  # What the heck?
>>> add.__name__  # The function is called "add"...
'wrapper'
>>> add.__doc__  # And where did the docstring go?!
>>>

The issue?
The issue is that applying the decorator @print_args replaced the function add with the function called wrapper that is defined inside print_args.
As I showed in the last issue, applying the decorator with the at sign @ is the same as this:

def add(a, b):
    ...

add = print_args(add)

So, the function add is being replaced by the function wrapper, but we do not want this to be visible.
The point of the decorator is to add some functionality in a way that is seamless...
Changing the name of the function and getting rid of its docstring does not seem like the way to go!

A partial solution

To fix this issue, you could be careful to copy things around inside your decorator.
For example, like this:

def print_args(f):
    def wrapper(*args, **kwargs):
        print(f"{args = } and {kwargs = }")
        return f(*args, **kwargs)

    wrapper.__name__ = f.__name__
    wrapper.__doc__ = f.__doc__
    return wrapper

This seems to fix some of the issues we have:

>>> @print_args
... def add(a, b):
...     """Adds two numbers together."""
...     return a + b
...
>>> add  # This still looks funky...
<function print_args..wrapper at 0x000001C58662B280>
>>> add.__name__  # but the name is correct...
'add'
>>> add.__doc__  # and so is this!
'Adds two numbers together.'

However, this was only part of the solution, because we see that the function add still looks weird in the REPL.
Not only that, but it is cumbersome to have to do all those assignments by hand every time you implement a decorator.
That is why the module functools has a tool called wraps.

The full fix

functools.wraps is a decorator (so meta, right?!) that fixes all of these things for you, so your decorated function preserves its name, its representation, its docstring, and some other things.
To use functools.wraps, use it as a decorator around the wrapper, and not the outer decorator itself:

def print_args(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        print(f"{args = } and {kwargs = }")
        return f(*args, **kwargs)
    return wrapper

Again, notice how wraps was used on the inner function wrapper, and not on the function print_args.
Then, you can use your decorator as usual:

>>> @print_args
... def add(a, b):
...     """Adds two numbers together."""
...     return a + b
...
>>> add
<function add at 0x000001C586709EE0>
>>> add.__name__
'add'
>>> add.__doc__
'Adds two numbers together.'

Why bother?

Why would you bother using functools.wraps in the first place, either way?
For several reasons, all related to the fact that Python has very powerful and useful introspection capabilities, and you probably do not want to mess with those.
For example, the built-in help accepts a function and gives you information about that function.
If you decorate your function and do not fix the name and the docstring, the built-in help won't work properly.

Got any decorator questions?

Like I mentioned already, on the 3rd of December I will be giving a remote workshop on decorators.
I will talk about defining your own decorators and functools.wraps, and I will also talk about decorators with arguments, other useful decorators available in the Standard Library, and more.

If you have any questions about decorators, send them in.
I will tailor the workshop to the sort of questions you send in, and I might even write one more newsletter issue on the subject of decorators, if you show interest!


Thanks for reading, and I'll see you next time!

Rodrigo.