A brief Python decorator primer

My last post looked at easing a common database idiom using a decorator. Interestingly, the brief discussion on reddit regarding the post had a number of people remarking that the coverage of decorators was more interesting than the SQL-related bits. Curious, that.

It seems I’m not alone in being occasionally puzzled by the way decorators work. It took me a little while to comprehend them. I’m still not a huge fan of how they work, but I can see their utility and do use them when appropriate. Almost without fail I end up heading over to PEP-318 and puzzling over the variations. The PEP isn’t terribly helpful since it is dominated by historical discussion rather than examples or instructions.

So let’s start with the basic no-argument decorator. It doesn’t take any declarative arguments when decorating a function; you just type @some_deco without parenthesis when using it.

def some_deco(f):
    def _inner(*args, **kwargs):
        print "Decorated!"
        return f(*args, **kwargs)

    return _inner

@some_deco
def some_func(a, b):
    return a + b

print some_func(1, 2)

This will print out:

Decorated!
3

According to the PEP, if you wrap some_func with @some_deco, it is the equivalent of some_func = some_deco(some_func). What it fails to call out is that this happens at load time. What we’re returning from the decorator function is a new function or callable that will be bound to the name some_func when this module is loaded. Subsequent run-time calls to some_func are actually calling the _inner function we defined on the fly when the decorator was called. Thanks to Python’s scoping rules and closures, we have access to the passed f parameter from the inner function, so we can call it.

Lovely. Now you have the power to tack on any sort of things you want to around a given function. Print a log message before or after, spawn a thread, add some exception handling. Whatever you want.

Things get a little more hairy if we want our decorators themselves to take parameters. According to the PEP, if was want parameters on the decorator to take parameters it’s the equivalent of func = decomaker(argA, argB, ...)(func). Again, this happens at load time, so it’s worth staring at for a moment. In plain English, this is saying that decomaker will be called with some arguments, and then the result of that will be called with a function reference. This is key! There are actually two invocations happening. The first will need to return a callable that accepts a function as it’s argument. That will then be called, and the result bound to the name func, so it better be callable as well. Let’s try it.

def log_wrap(message):
    """Print `message` each time the decorated function is called."""
    def _second(f):
        def _inner(*args, **kwargs):
            print message
            return f(*args, **kwargs)

        # This is called second, so return a callable that takes arguments
        # like the first example.
        return _inner

    # This is called first, so we want to pass back our inner function that accepts
    # a function as argument
    return _second

@log_wrap("Called it!")
def func(a, b):
    return a + b

print func(1, 2)

This will output:

Called it!
3

Ok, ok, that’s not terribly useful, but it illustrates the point (and expands on the first example by being a parameterized version). It requires that we construct two callable items inside our decorator, and that each will be called in turn during load time. Let’s quick dissect what happens at load time. When python gets to the declaration of func with the decorator sitting on top of it, it calls the decorator function with its argument (in this case, the string “Called!”). What is returned is a function reference to _second. Next, that just-returned reference to _second is invoked with a single argument: a function pointer to func, which is what we’re actually trying to decorate. The return value of _second is a reference to _inner, which will be bound to the name func, ready for run-time invocation.

Yeah, I’m a little dizzy, too. But if you practice this a couple times, it’ll start to make sense. Again, the key to getting this second form working is to recognize that two calls are happening at load time and returning the appropriate callable references at the right time.

Once you have it down, use this new hammer when it serves you, remembering that not every problem is a nail.