How to Define Function Decorators with Functools.wraps in Python 3

548
SHARES
2.5k
VIEWS

Python has special syntax for decorators that can be applied to functions. A decorator has the ability to run additional code before and after each call to a function it wraps. This means decorators can access and modify input arguments, return values, and raised exceptions. This functionality can be useful for enforcing semantics, debugging, registering functions, and more.

For example, say that I want to print the arguments and return value of a function call. This can be especially helpful when debugging the stack of nested function calls from a recursive function. Here, I define such a decorator by using *args and **kwargs to pass through all parameters to the wrapped function:

def trace(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        print(f'{func.__name__}({args!r}, {kwargs!r}) '
              f'-> {result!r}')
        return result
    return wrapper

I can apply this decorator to a function by using the @ symbol:

@trace
def fibonacci(n):
    """Return the n-th Fibonacci number"""
    if n in (0, 1):
        return n
    return (fibonacci(n - 2) + fibonacci(n - 1))

Using the @ symbol is equivalent to calling the decorator on the function it wraps and assigning the return value to the original name in the same scope:

fibonacci = trace(fibonacci)

The decorated function runs the wrapper code before and after fibonacci runs. It prints the arguments and return value at each level in the recursive stack:

fibonacci(4)

>>>
fibonacci((0,), {}) -> 0
fibonacci((1,), {}) -> 1
fibonacci((2,), {}) -> 1
fibonacci((1,), {}) -> 1
fibonacci((0,), {}) -> 0
fibonacci((1,), {}) -> 1
fibonacci((2,), {}) -> 1
fibonacci((3,), {}) -> 2
fibonacci((4,), {}) -> 3

This works well, but it has an unintended side effect. The value returned by the decorator—the function that’s called above—doesn’t think it’s named fibonacci:

print(fibonacci)

>>>
.wrapper at 0x108955dc0>

The cause of this isn’t hard to see. The trace function returns the wrapper defined within its body. The wrapper function is what’s assigned to the fibonacci name in the containing module because of the decorator. This behavior is problematic because it undermines tools that do introspection, such as debuggers .

For example, the help built-in function is useless when called on the decorated fibonacci function. It should instead print out the docstring defined above ('Return the n-th Fibonacci number'):

help(fibonacci)

>>>
Help on function wrapper in module __main__:

wrapper(*args, **kwargs)

Object serializers break because they can’t determine the location of the original function that was decorated:

import pickle

pickle.dumps(fibonacci)

>>>
Traceback ...
AttributeError: Can't pickle local object 'trace..
➥wrapper'

The solution is to use the wraps helper function from the functools built-in module. This is a decorator that helps you write decorators. When you apply it to the wrapper function, it copies all of the important metadata about the inner function to the outer function:

from functools import wraps

def trace(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
         ...
    return wrapper
@trace
def fibonacci(n):
    ...

Now, running the help function produces the expected result, even though the function is decorated:

help(fibonacci)

>>>
Help on function fibonacci in module __main__:
fibonacci(n)
      Return the n-th Fibonacci number

Leave a Reply

Your email address will not be published. Required fields are marked *