2019-10-22T15:32:45Z

The Ultimate Guide to Python Decorators, Part II: Altering Function Behavior

Welcome to the second part of my Python decorator series. In the first part I showed you the most basic style of decorators, which are used to register functions as handlers or callbacks for events. In this part I'm going to show you more interesting decorators that alter or complement the behavior of the decorated function.

Simple Decorators Review

Before we delve into new territory, let's review how the simple decorators from Part I work. Below you can see one, which I entered into a Python shell (feel free to start a shell and type it if you want to play along with me!):

>>> def my_decorator(f):
...     print('decorating', f)
...     return f
...
>>>

This decorator example is perfectly valid, but to keep it simple it does nothing besides printing a message to the console. If I wanted to use this decorator on a function, I could do something like this:

>>> @my_decorator
... def my_function(a, b):
...     return a + b
...
decorating <function my_function at 0x10ae241e0>
>>> my_function(1, 2)
3
>>>

Note how the decorating ... print statement appears when the function is defined, not when it is invoked. This is because Python calls the decorator function at the time the decorated function is declared. This decorator doesn't do anything, so obviously my_function() is unaffected by it and can be called as usual.

Decorators That Replace The Decorated Function

The more advanced decorators that I'm going to show you next are not going to return the same function like the one above, they are going to return a different function, and this is going to enable an array of very cool tricks the simpler decorators cannot do. Almost always, these decorators are implemented with inner functions, which for many developers are an odd and obscure feature of the Python language. To introduce them to you gently, I'm going to convert the above "do-nothing" decorator into the more powerful style, but I'm going to do it in small incremental steps.

Let's begin by changing the return statement to return a function other than f, once again in a Python shell:

>>> def forty_two():
...     return 42
...
>>> def my_decorator(f):
...     print('decorating', f)
...     return forty_two
...
>>>

Here I defined a function forty_two(), and then I had my decorator return a reference to that function instead of the f that I used above. Let's use this decorator as I did above to understand what this change does:

>>> @my_decorator
... def my_function(a, b):
...     return a + b
...
decorating <function my_function at 0x10ae24268>
>>> my_function(1, 2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: forty_two() takes 0 positional arguments but 2 were given
>>>

So what happened here? Somehow the decorated function appears to be broken now and can't be called as I did in the previous example. Recall that the function returned by the decorator function replaces the original function. The decorator is basically changing my_function to be a reference to forty_two, so when I called my_function(1, 2) I got an error, because I'm really calling forty_two(1, 2) which is invalid, as this function takes no arguments.

Note that the error message names the function as forty_two, even though I defined it with the name my_function. This little naming side effect of decorators that return a different function can be very confusing and has a solution, but that is a topic for a future article in the series.

How do you think the error can be resolved? Since my_function is now an alias to forty_two, I have to call this function with no arguments:

>>> my_function()
42

You may be wondering where is the original function now. In this example, sadly the original function is gone, since it was replaced by the forty_two() function by the decorator. So obviously this is a terrible decorator that has absolutely no use. The idea is that the returned function should act as a wrapper for the original function, not as a complete replacement. And this is when the concept of inner functions becomes relevant. Take a look at the next decorator:

>>> def my_decorator(f):
...     def wrapped():
...         return f(1, 2)
...     print('decorating', f)
...     return wrapped
...
>>>

Here the wrapped() function is said to be an inner function because it is defined inside the body of another function. The big feature of inner functions is that they carry with them any variables in the scope of the parent function, what in computer science is called a closure. Note how inside the body of wrapped() I can call f(1, 2), even though f isn't defined within the function and instead is an argument of the parent my_decorator() function.

Let's try this new decorator:

>>> @my_decorator
... def my_function(a, b):
...     return a + b
...
decorating <function my_function at 0x10ae24378>
>>> my_function(1, 2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: wrapped() takes 0 positional arguments but 2 were given
>>> my_function()
3
>>>

Interesting, right? I stil can't pass the a and b arguments of the original my_function(), but when I call it with no arguments I get 3 as a result, because the wrapped() function injects the arguments 1 and 2 on its own.

To fully reinstate the two arguments from the original function, I can define wrapped() also with two arguments, and pass those arguments through to f:

>>> def my_decorator(f):
...     def wrapped(a, b):
...         return f(a, b)
...     print('decorating', f)
...     return wrapped
...
>>>

But of course, this decorator can only be applied to functions that take two arguments. To make wrapped() work as a wrapper for all functions, the best solution is to code it so that it accepts any arguments, and then passes them on to f. In Python, a function can be made to accept any arguments with the *args and **kwargs special arguments, which represent positional and keyword arguments respectively.

>>> def my_decorator(f):
...     def wrapped(*args, **kwargs):
...         return f(*args, **kwargs)
...     print('decorating', f)
...     return wrapped
...
>>>

And this is finally a "do-nothing" decorator that is compatible with all functions regardless of their arguments, implemented with an inner function.

A decorator structured in this way can insert additional behaviors that enhance what the decorated function does, either before or after that function is called inside wrapped(). It can also decide to not invoke the decorated function in certain cases, for example if it determines that the arguments that were passed are incorrect.

See the comments in the following example, which show where in the decorator these new behaviors can be inserted:

def my_decorator(f):
    def wrapped(*args, **kwargs):
        # ...
        # insert code that runs before the decorated function
        # (and optionally decide to not call that function)
        # ...
        response = f(*args, **kwargs)
        # ...
        # insert code that runs after the decorated function
        # (and optionally decide to change the response)
        # ...
        return response
    return wrapped

To help you visualize how this decorator style works, here is an example that uses print statements in all the significant parts:

>>> def my_decorator(f):
...     def wrapped(*args, **kwargs):
...         print('before function')
...         response = f(*args, **kwargs)
...         print('after function')
...     print('decorating', f)
...     return wrapped
...
>>> @my_decorator
... def my_function(a, b):
...     print('in function')
...     return a + b
...
decorating <function my_function at 0x10ae24268>
>>> my_function(1, 2)
before function
in function
after function

Example #1: Injecting New Arguments

An obvious trick that can be implemented with this type of decorators that I hinted at above is to change the argument list that is sent to the decorated function when it is invoked. As an example, consider the following decorator, which injects the current time as a first argument into any functions it decorates:

>>> from datetime import datetime
>>> def add_current_time(f):
...     def wrapped(*args, **kwargs):
...         return f(datetime.utcnow(), *args, **kwargs)
...     return wrapped

And here is an example usage:

>>> @add_current_time
... def test(time, a, b):
...     print('I received arguments', a, b, 'at', time)
...
>>> test(1, 2)
I received arguments 1 2 at 2019-10-10 21:38:35.582887

As you can see, the decorated function is written to accept a first time argument, but this argument is automatically added by the decorator, so the function is invoked with the remaining arguments, in this case a and b.

Example #2: Altering the Return Value of a Function

Another very common task decorators are good for is to perform a conversion on the return value of the decorated function. If you have many functions that invoke a data conversion function before returning, you can move this conversion task to a decorator to make the code in your functions simpler, less repetitive and easier to read.

A good example of this technique can be applied to my favorite web framework, Flask. Consider the following Flask view function:

@app.route('/')
def index():
    return jsonify({'hello': 'world'})

If your Flask application implements an API, you likely have many routes that end by returning a JSON payload, generated with the jsonify() function. Wouldn't it be nice if you could just return the dictionary, instead of having to include the jsonify() call at the end of every view function? A to_json decorator can perform the conversion to JSON for you:

@app.route('/')
@to_json
def index():
    return {'hello': 'world'}

Here the index() function returns the dictionary, which for Flask is an invalid response type. But the to_json decorator wraps the function, and has a chance to perform the conversion and update the response before it reaches Flask. Here is the complete application, including the implementation of this decorator:

from flask import Flask, jsonify

app = Flask(__name__)

def to_json(f):
    def wrapped(*args, **kwargs):
        response = f(*args, **kwargs)
        if isinstance(response, (dict, list)):
            response = jsonify(response)
        return response
    return wrapped

@app.route('/')
@to_json
def index():
    return {'hello': 'world'}

The wrapped() function simply invokes the original function, called f in this context, and then checks if the return value is a dictionary or list, in which case it calls jsonify(), effectively intercepting and fixing the return value before it is returned to the framework.

Note: As of Flask version 1.1 you can, in fact, return a dictionary, and Flask automatically converts it to JSON. So while the above decorator is not needed anymore, it is still a good example of building filters using the decorator pattern.

Example #3: Validation

Yet another useful technique that can be implemented with decorators is to perform any kind of validation before the decorated function is allowed to run. A very common example in a web application is to authenticate the user. If the validation/authentication task ends in a failure, then the decorated function is not invoked, and instead the decorator raises an error.

Here is an example of this technique for Flask:

>>> from flask import request, abort
>>> ADMIN_TOKEN='fheje3$93m*fe!'
>>> def only_admins(f):
...     def wrapped(*args, **kwargs):
...         token = request.headers.get('X-Auth-Token')
...         if token != ADMIN_TOKEN:
...             abort(401)  # not authorized
...         return f(*args, **kwargs)
...     return wrapped
...
>>> @app.route('/admin')
... @only_admins
... def admin_route():
...     return "only admins can access this route!"
...

In this example the only_admins decorator looks for a X-Auth-Token header in the incoming request, and then checks that it matches a secret administrator token, that for simplicity I have set as a constant. If the token header isn't present, or if it is present but the token does not match, then the abort() function from Flask is used to generate a 401 response and halt the request. Otherwise the request is allowed to go through by invoking the decorated function.

Note how the admin_route() example view function has both the app.route and only_admins decorators applied. This is called decorator chaining. Chaining decorators is a tricky topic so I'm going to dedicate some time to it in a future article. For those of you working with Flask, assume that almost always the app.route decorator will be the first decorator in a chain.

Decorators with Arguments

So far I have only shown you decorators that do not take any arguments themselves, the arguments always go to the decorated function. I have done this on purpose, because decorators that accept arguments of their own in addition to those that are passed through to the decorated function are more complex to write. This is going to be the topic of the next articles, so make sure you familiarize yourself with the concepts I covered in this and the previous article and be ready for another jump in complexity in the next installment!

2 comments

  • #1 Brandon said 2019-10-29T14:49:26Z

    Very helpful, Miguel! Thanks!

  • #2 Charley said 2019-11-10T20:05:09Z

    Thank you for this post. I look forward to future installments!

Leave a Comment