Decorators deep dive

I thought I’d do a post here getting incredibly painfully specific about how decorators work, since when I was learning about them early on I felt this was missing. Note Python decorators are distinct from, though related to, the decorator pattern (which is also implementable in Python).

Preliminaries

This post assumes knowledge of basic Python, as well as some of its implementation of typing. Some general (but very basic) type theory is also helpful, mainly because I’m not too careful below to distinguish between pythonic type notation and more general abstract typing (the difference between f: Callable[[T], S] and f: T -> S which mean the same thing, that being that f is a function which maps Ts to Ss). This article will focus on decorators as applied to functions, though very similar principles apply when decorating classes.

Basic definition

The first thing to note about decorators is their appearance as @decorator is just ‘syntactic sugar’, ie syntax chosen to make things easier to read. All that

@decorator
def f(x):
	return x

really does is expand to

def f(x):
    return x
f = decorator(f)

so that in the future when you call f(x), you’re really calling [decorator(f)](x). (As you’ll see in this article, I’m a big fan of bracketing to help make function composition clearer – this bias comes from suffering through a lot of painful category theory proofs at uni.)

Now the above is cool, but doesn’t actually show us how decorators can be implemented – how we can build our own. What we can learn from the above that if f has signature T -> S (ie it takes in things of type T and returns things of type S), decorator(f) must be a function of (almost) the same signature. In turn what this tells us about decorator is that it must be a function which takes in functions of type T -> S and returns a function of type T -> R (for some type R which may or may not be the same as S).

To be a bit more concrete, if we have f: int -> int, for example

def f(i: int) -> int:
    return i + 1

then we must have decorator(f): int -> R (and in all likelihood we would in fact expect R = int – a simplification we will make going forwards to keep the typing explicit), which in turn tells us decorator: Callable[[int], int] -> Callable[[int], int] (or, to be rigidly pythonic, decorator: Callable[[Callable[[int], int]], Callable[[int], int]]). We can leverage this insight to help us build a functioning implementation of a decorator.

What we note, in line with the typing we derived above, is that the decorator must be of the form

def decorator(func: Callable[[int], int]) -> Callable[[int], int]:
    [innards]
    return some_func

where some_func: int -> int. Well the only place we can define some_func is inside the decorator, so we arrive at an instantiation of the standard decorator definition (with our choice of typing). I’ve chosen to call some_func the more standard inner, and I’ve also given it a specific definition, so we can work through the definition below explicitly.

def decorator(func: Callable[[int], int]) -> Callable[[int], int]:
    def inner(i: int) -> int:
        return func(i) + 1
    return inner

Using what we noted above about syntactic sugar, we can now check explicitly what our decorator does. We know that

@decorator
def f(i: int) -> int:
    return i + 1

really means

def f(i: int) -> int:
    return i + 1
f = decorator(f)

So now when in the main body of our code we call f(i), really we call [decorator(f)](i), which is inner(i), which is f(i) + 1 (where in this last statement, f refers to the ‘old’ f, before it has been redefined, because we are evaluating code inside the decorator), which finally is (i + 1) + 1 = i + 2.

Note our choice of an f which takes one integer argument and returns an integer was just to make the examples simple, the above extends to a more general definition (dropping explicit typing for ease).

def more_general_decorator(func):
    def inner(*args, **kwargs):
        [maybe do some stuff]
        return func(*args, **kwargs)
    return inner


@more_general_decorator
def g(*args, **kwargs):
    [do some stuff]

Note that inner doesn’t need to return func(*args, **kwargs) or in fact anything at all – there are loads of things it could do depending on what we want from the decorator – but this is a common choice.

Decorators with arguments

What about if we wanted our decorator to be able to take arguments, such as with the flaky decorator, used to make unreliable tests automatically rerun on failure to a fixed limit? This has arguments max_runs, min_passes, and rerun_filter, for example. Once we note the basic syntactic sugar rule still applies, so

@decorator_with_args(y)
def f(x):
    return x

really expands to f = [decorator_with_args(y)](f), we see that all that needs to be the case is that decorator_with_args(y), instead of just decorator as in the above, now needs to be the mapping from f to a new function. Now decorator_with_args is really a decorator factory – it is a mapping from arguments (of some type we haven’t specified) to decorators.

With this insight we can see its definition is going to take the form (we’ll ditch typing for now as I don’t think it adds much here)

def decorator_with_args(y):
    [do some stuff]
    return decorator

So it’s clear all we need to do is construct a decorator in the normal way inside decorator_with_args:

def decorator_with_args(y):
    def decorator(func):
        def inner(*args, **kwargs):
            [do some stuff]
	        return func(*args, **kwargs)
	    return inner
	return decorator

To take a concrete example, building on our toy functions above, lets say we have f be the successor function as above. Now define

def add_to_output(extra: int) -> Callable[[Callable[[int], int]], Callable[[int], int]]:
    def inner_decorator(func: Callable[[int], int]) -> Callable[[int], int]:
        def inner(i: int) -> int:
            return func(i) + extra
        return inner
    return inner_decorator

Now if we define

@add_to_output(3)
def f(i: int) -> int:
    return i + 1

Then expanding, we have

f(i) = [[add_to_output(3)](f)](i)
	 = [inner_decorator(f)](i)
	 = [inner(i)]
	 = f(i) + 3
	 = (i + 1) + 3
	 = i + 4

Here’s another fun example I used in code a while back, to make an arbitrary function (in this case a test) run n times:

def repeat(repetitions):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(repetitions):
                func(*args, **kwargs)
        return wrapper
    return decorator

Decorators with optional arguments

So we can define decorators which accept args – cool! What about if we wanted our decorator to take arguments optionally? There is a naïve solution building on the above where we simply define the decorator as above but with some default values in the definition – lets take my repeat decorator and have it default to 5 repetitions

def repeat(repetitions=5):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(repetitions):
                func(*args, **kwargs)
        return wrapper
    return decorator

We could then use it as follows:

@repeat()
def f() -> None:
    print('Hello!')

and we’d get Hello! printed 5 times in the output. This is useful, but doesn’t feel optimal, since we have to pass the empty brackets after @repeat in the definition of f. Of course it wouldn’t work if we did things like this as stands, because repeat is a decorator factory, and not a decorator itself, so when we define

@repeat
def f() -> None:
    print('Hello!')

We would have f = repeat(f) which breaks because f is taken as the argument of repeat (which should be either nothing, or an integer representing the number of repeats), and not passed to decorator, which then is missing an argument. Indeed, trying to run the above returns

E TypeError: repeat.<locals>.decorator() missing 1 required positional argument: 'func'

But all is not lost! We can in fact define a decorator which can do both, with a little bit of argument inspection. Essentially we need repeat to function as a decorator factory if it gets arguments, and just as a decorator if not. In code this looks like

def repeat(*args):
    if len(args) == 1 and isinstance(args[0], Callable):
        def wrapper(*func_args):
            func = args[0]
            for _ in range(5):
                func(*func_args)
        return wrapper
    else:
        def decorator(func):
            def other_wrapper(*func_args):
                repetitions = args[0]
                for _ in range(repetitions):
                    func(*func_args)
            return other_wrapper
        return decorator

What we do here is we check whether the decorator has received one argument which is a function, in which case we are in the ’no arguments’ case, and act accordingly, else we are in the ‘arguments case’. The above makes structural assumptions about the arguments fed to the decorator, and isn’t totally explicit, so if using this in production, it is probably sensible to add some comments. The implementation here also doesn’t include kwargs, however these can be added in an analogous manner.

Note here is a case where the assumption that we’re decorating a function is important – if we were to decorate a class with the above, unless the class was a sub-type of Callable, the first check would always fail.

A real-world example of the above pattern (or a variation on it) can be seen in the flaky operator, mentioned above: this optionally takes the key-word arguments max_runs, min_passes, and rerun_filter.