Creating and Using Decorators

Decorators are very useful, powerful tools and, once properly understood, easy to implement. To be sure, this is what we are talking about:

@decorator
def func():
    pass

If you have seen this syntax and want to know what it means and how to use it, then keep reading.

What is a decorator? Here we see a decorator named decorator being attached to a func() function, whence it follows that decorators are something that is attached to functions. Right! In fact, decorators are themselves functions, that take a function as an argument and return another function. Furthermore, the previous code is just syntactic sugar for:

def func():
    pass
func = decorator(func)

Now the proposed definition makes a bit more sense: Our decorator is a function to which we pass another function as an argument (func()) and which returns a new function that gets assigned to func. Of course, we haven't yet defined the code for this to actually happen.

The reason why the language includes the @ syntax since Python 2.4 is mainly because calling decorators as we did in our second code can leave the decorator very far from the function definition if its body is very extensive. On the contrary, the @ syntax allows us to clearly see which decorators are being attached to a function in its very definition.

Let's take the following function as an example.

def add(a, b):
    return a + b

Suppose that we want to print a message every time this function gets called. We could achieve this by using a decorator. Let's call it debug.

@debug
def add(a, b):
    return a + b

Now it's time to write its code. Remember that a decorator must be a function that takes another function as an argument.

def debug(f):
    pass

But, in addition, it has to return a new function. So we'll do exactly that.

def debug(f):
    def new_function():
        pass
    return new_function

This is completely valid! Python allows us to define functions within functions. And we said that the code would be equivalent to:

def add(a, b):
    return a + b
add = debug(add)

For every call to add() in our program to continue to work, the decorator must return a new function with the same argument structure and behavior as the original function (that is, take two objects as arguments and return the sum that results from them). However, our new_function() doesn't take any arguments. Let's fix that.

def debug(f):
    def new_function(a, b):
        pass
    return new_function

Finally, we don't need to replicate the add() code in new_function(), since we already have a reference to the former in the f argument. We can simply do:

def debug(f):
    def new_function(a, b):
        return f(a, b)
    return new_function

Great! The decorator can now be attached to our function without errors. Let's finally print a message when the function gets called (that was the whole point of the decorator!).

def debug(f):
    def new_function(a, b):
        print("Function add() called!")
        return f(a, b)
    return new_function
@debug
def add(a, b):
    return a + b
# A message will be printed on every call to add().
print(add(7, 5))

Congratulations! That was our first experiment with decorators.

Let's make a further step. Consider this other function that returns the opposite of a number n.

def neg(n):
    return n * -1

Since our decorator was created specifically for the add() function, which required two arguments, it will fail when attached to neg().

@debug
def neg(n):
    return n * -1
#TypeError!
print(neg(5))

This throws:

TypeError: new_function() missing 1 required positional argument: 'b'

The solution would be to adjust the structure of new_function() to take a single argument (such as neg()). But that would cause the decorator to fail when attached to add(). A better approach is to include a more generic notation to accept arbitrary positional and keyword arguments:

def debug(f):
    def new_function(*args, **kwargs):
        print(f"Function {f.__name__}() called!")
        return f(*args, **kwargs)
    return new_function

(Note that the f in the third line means string formatting, while the f in the fourth line is a reference to the decorated function, such as add() or neg().)

We also changed the message in order to always include the name of the decorated function (f.__name__). Now our decorator is ready to be applied to any function regardless its argument structure.

Decorators With Arguments

Even more interesting, decorators might also have arguments. Suppose we want our debug decorator to include an option to pause the program execution and launch the debugger when the decorator function gets called. The syntax could read like this:

@debug(breakpoint=True)
def func():
    pass

How could we implement this? Let's translate the code to its equivalent without the @ syntax, as we have seen.

def func():
    pass
func = debug(breakpoint=True)(func)

Now it is clear that the argument of the debug() function cannot be f, but rather breakpoint, and that it must return our previous decorator, to which the func() function will be passed. It would be something like this:

def debug(breakpoint=False):
    def debug_decorator(f):
        def new_function(*args, **kwargs):
            print(f"Function {f.__name__}() called!")
            return f(*args, **kwargs)
        return new_function
    return debug_decorator

As you can see, we create a new function debug() that will be in charge of receiving the arguments and we define our decorator within it as debug_decorator.

Lastly, we should include the code to start the debugger where appropriate by using the pdb standard module.

import pdb
def debug(breakpoint=False):
    def debug_decorator(f):
        def new_function(*args, **kwargs):
            print(f"Function {f.__name__}() called!")
            if breakpoint:
                pdb.set_trace()
            return f(*args, **kwargs)
        return new_function
    return debug_decorator

And let's do some tests:

@debug(breakpoint=True)
def add(a, b):
    return a + b
@debug()  # Parentheses required!
def neg(n):
    return n * -1
print(neg(5))
print(add(7, 5))  # This call starts the debugger.

Further Considerations

A common problem when using decorators is that since they encapsulate the function they are attached to, we can no longer access the original function's attributes, such as its name (__name__) or documentation (__doc__). Consider the following code.

def debug(f):
    def new_function(*args, **kwargs):
        print(f"Function {f.__name__}() called!")
        return f(*args, **kwargs)
    return new_function
@debug
def neg(n):
    "Return the inverse of n."
    return n * -1
print(neg.__name__)  # new_function
help(neg)

We see that the __name__ and __doc__ attributes (accessed by help()) return the information of new_function(), not neg(). To avoid this, we can use the standard decorator functools.wraps(), which is responsible for copying all the attributes of the original decorated function to the new one.

from functools import wraps
def debug(f):
    @wraps(f)
    def new_function(*args, **kwargs):
        print(f"Function {f.__name__}() called!")
        return f(*args, **kwargs)
    return new_function

As simple as that!