Python decorators with optional arguments

Ever seen decorators which can be used like this?

@baked
def get_cake(flavour):
    # …

@baked(temperature=180, duration=25)
def get_cake(flavour):
    # …

Django’s template filter registration decorator is a good example of this, where it can be called as either a decorator, or a function that returns a decorator (specifically, a function that returns a function that takes a function and returns a function :D).

All these levels of indirection can get a little confusing. First, lets look at a simple decorator function:

import functools

def baked(method):
    @functools.wraps(method)
    def f(*args, **kwargs):
        thing_to_be_baked = method(*args, **kwargs)
        return bake(thing_to_be_baked)
    return f

functools.partial is a useful utility function that copies attributes from the wrapped function to the wrapping function.

Next, here’s one that takes additional arguments:

import functools

def baked(temperature=None, duration=None):
    def decorator(method):
        @functools.wraps(method)
        def f(*args, **kwargs):
            thing_to_be_baked = method(*args, **kwargs)
            return bake(thing_to_be_baked, temperature, duration)
        return f
    return decorator

Here, baked is the function that returns a function (decorator) that takes a function (method) and returns another function (f). This is a lot of nesting, and still doesn’t handle the case where the user doesn’t want to supply the optional arguments.

We can reduce the nesting using functools.partial, while at the same time making the arguments optional:

import functools

def baked(method=None, temperature=None, duration=None):
    # If called without method, we've been called with optional arguments.
    # We return a decorator with the optional arguments filled in.
    # Next time round we'll be decorating method.
    if method is None:
        return functools.partial(baked, temperature=temperature, duration=duration)
    @functools.wraps(method)
    def f(*args, **kwargs):
        thing_to_be_baked = method(*args, **kwargs)
        return bake(thing_to_be_baked)
    return f
Posted in Python | Leave a comment

Leave a Reply