Composing Decorators in Python

The original version of this post is available on my website.

While implementing a PEG parser, my friend and I stumbled upon a weird-looking yet valid Python program. It looked something like this:

def foo():
    print("Hello, World!")


def bar():
    return foo


def baz():
    return bar


baz()()()

What we have here is something similar to a Closure. It's a way to use first-class functions to somewhat "bind" a function and an environment together. After a bit of thought, I realized that this is very similar to Python's decorators.

Python's @decorators

Here's an example of a function decorator in action:

def stringify(f):
    """
    Decorator that converts the output of a function to a string.
    """

    def inner(x):
        return str(f(x))

    return inner


@stringify
def square(x):
    """
    Returns the square of x as a string.
    """
    return x * x


result = square(4)

>>> result
'16'

In the snippet above, we use the stringify decorator to convert all output of the square function to a string.

This seems like a helpful way to organize code, but this got me wondering - can I stack multiple decorators? Turns out you can! Here's an example:

def double(f):
    """
    Decorator that doubles the output of a function.
    """

    def inner(x):
        return f(x) * 2

    return inner


def stringify(f):
    """
    Decorator that converts the output of a function to a string.
    """

    def inner(x):
        return str(f(x))

    return inner


@double
@stringify
def square(x):
    """
    Returns the square of x as a string.
    """
    return x * x


result = square(4)

>>> result
'1616'

With this new double decorator applied, we repeat the string '16' twice - getting '1616' as the output. We can also switch the order of the decorators around to get '32' as the output. Python's decorator's are applied closest-to-function first.

But.... what if I want to decorate my decorators?

Decorating the decorators

Knowing I can stack decorators, I wanted to make decorators for decorators. This ended up looking something like this:

def stringify(wrapper):
    # we need to get the function f, and the input x
    def s(f):
        f = wrapper(f)

        def t(x):
            result = str(f(x))
            return result

        return t

    return s

@stringify
def double(f):
    def d(x):
        return f(x) * 2

    return d

This seems like a lot of work just to decorate decorators - what if I wanted to decorate the stringify decorator? These decorators that are supposed to help soon become lots of boilerplate to manage.

Composition to the rescue

One of my favourite features from Haskell and functional programming is function composition. It's a simple way to "pipeline the result of one function, to the input of another, creating an entirely new function".

There's a solution by using the functools.reduce function to create a simple compose function to easily compose multiple functions down to a decorator.

@compose(increment, double, stringify)
def foo(x):
    return x * x

While this optimally solves the problem, the syntax just doesn't look quite as appealing as the Haskell version. So instead, I implemented this simple class to let me use the + operator to compose together a decorator.

class Compose:
    def __init__(self, f=None):
        self.f = f if f else lambda x: x
    def __add__(self, other):
        return Compose(lambda x: other(self.f(x)))
    def __call__(self, f):
        return lambda x: self.f(f(x))

decorate = Compose()

def double(x):
    return x * 2

def stringify(x):
    return str(x)

def increment(x):
    return x + 1

@(decorate + increment + double + stringify)
def foo(x):
    return x * x

Much better.

Wrapping up

Decorators are a yet another part of the pile of syntactic sugar in Python. This post covers some niche use-cases for them, and how using some functional programming concepts such as composition can help us write better code.

If you want to learn more about decorators, there's an amazing Primer on Python Decorators online.

If you liked reading this, check out my website at krishkrish.com for more.