An introduction to Python decorators

May 11, 2022

I’d never taken the time to wrap my head around decorators in Python, but I watched a video from PyCon 2021 today which explained them well. I decided to write this post to share the knowledge; read on for my explanation of decorators. In this post, I assume the reader is comfortable with functions in Python. If you’re not, you might want to do some preliminary reading.

Intro: passing functions as arguments

You probably know that you can pass functions as arguments to other functions. This is often used in e.g. the map function. For example, we can define a function double that doubles its argument:

def double(x):
    return x*2

If we check in the REPL, we see that double is a function:

>>> double
<function double at 0x10638de50>

And we can pass it to map and have it modify a list of numbers:

>>> list(map(double, [1,2,3]))
[2, 4, 6]

It’s good to keep this in mind as we get started with decorators.

A basic decorator with no arguments

Let’s say we want to create a decorator which ensures that a function’s result will always be even. We will have some function func, which will generate a random number, and we want to keep calling it until the result is even.

A higher-level function that would provide this guarantee for a func could look like this:

def even_only(func, *args, **kwargs):
    while True:
        result = func(*args, **kwargs)
        if result % 2 == 0:
            return result

Nothing special to say here, if you know a bit of Python, this is pretty self-explanatory.

Now, to actually use this, we need a function to generate a random number. That function could be anything, but for our purpose let’s just use a call to random.randint:

import random
def get_number():
    return random.randint(0, 101)

Now if we want to use this function in conjunction with even_only, we call it as such:

>>> even_only(get_number)
30

The problem here is that we always need to include the call to even_only, which means a lot of typing. Also, if we forget to include a call to even_only, we might get an odd number, and we always want an even number. That is, the function get_number itself is not guaranteed to return an even number. So, let’s convert even_only to a decorator and get that guarantee.

Here’s how we’d modify even_only to work as a decorator:

def even_only(func):
    def wrap(*args, **kwargs):
        while True:
            result = func(*args, **kwargs)
            if result % 2 == 0:
                return result
    return wrap

The decorator even_only “wraps” the original function. The only new lines here are def wrap(*args, **kwargs), which defines the wrap function and allows any arguments to be passed into the function being wrapped (get_number in this case), and return wrap which returns the function.

Why does even_only have to return the wrap function (or in general, why does it need to return a function)? Well, here’s how you’d use even_only as a decorator for get_number:

import random

@even_only
def get_number():
    return random.randint(0, 101)

The @ syntax means that the definition of get_number is replaced by the result of even_only(get_number). It’s as if you wrote:

import random
def get_number():
    return random.randint(0, 101)

get_number = even_only(get_number)

You then want to call get_number() to get a random number; this is why even_only has to return a function (in fact, every function decorator has to return a function).

Now if you call the decorated get_number, you can be sure that it’ll always return an even number:

>>> for _ in range(200):
...     assert(get_number() % 2 == 0)
...
>>>

A decorator with arguments

Decorators can also take arguments. Let’s say instead of even numbers, we want to ensure that the generator function will always return a number divisible by some n.

Here’s how we’d create such a decorator, let’s call it only_divisible_by:

def only_divisible_by(n):
    def decorate(func):
        def wrap(*args, **kwargs):
            while True:
                result = func(*args, **kwargs)
                if result % n == 0:
                    return result
        return wrap
    return decorate

The internals of the function are the same as before, except we have to add one extra level of function definition, because only_divisible_by is now a function itself, which accepts a single argument. Let’s see how that looks in practice and clarify why you need that extra function. This is how you’d use the decorator:

@only_divisible_by(3)
def get_number():
    return random.randint(0, 101)

Let’s unroll the decorator syntax to see why we need three levels of functions:

def get_number():
    return random.randint(0, 101)

get_number = (only_divisible_by(3))(get_number)

You see here that the first function call, only_divisible_by(3), must evaluate to a function itself, which we can interpret as an ‘instance’ of the inner decorate function with n set to 3. We then call that instance of decorate with the function get_number as an argument, which must then return a function that we can call to generate numbers (get_number()). So in the end, get_number is a call to the inner wrap function with n set to 3 and func set to the original get_number.

When using this decorated function, we know that the result will always be divisible by whatever n we choose (3 in this case):

>>> for _ in range(200):
...     assert(get_number() % 3 == 0)
...
>>>