Sunday, June 5, 2011

Python Decorator with Optional Keyword Arguments


Decorators in Python are extremely useful tools. This syntactic sugar is an integral part of standard python code, and can be custom-built to make your own code more succinct and awesome.

I've used decorators to mark functions as deprecated, make them curryable, make sure a user is authenticated (when that method was a render method in a Twisted resource), and more. Recently I was using them to add a little bit of magic to a collection of methods that I needed to use a lot, that had a lot of repeated code in them. I needed to sometimes set up the executing environment of those functions a little differently than usual though, using keyword options that I would pass to the decorator. I've seen decorators used this way before, so I figured it would be easy. Uhm, not so much.
Decorators are a little confusing to people who haven't written them before. The idea is simple:
@decorator
def function():
    pass
 
def function():
    pass
function = decorator(function)
A decorator is basically just a function (it can also be done with a class) which takes a function as its argument, and returns a function. The decorator in the example above might have been:
def decorator(func):
    def inner(*args, **kwargs):
        print "This is a decorated function!"
        func(*args, **kwargs)
    return inner
 
@decorator
def function(x):
    print x
 
function(42)
#> This is a decorated function!
#> 42
The decorator takes a function. It creates a new, "inner" function and closes the passed function. In the inner function, we print something, then call the original function. You can see this in action in this example, where calling the decorated function results in the text "This is a decorated function!" being printed.
A decorator can also take arguments:
@decorator(argument="An argument")
def function(x):
    print x
Remembering that decorators are just syntactic sugar, and referring to the first example I gave, what do you think the equivalent statement is? I figured it must be something like function = decorator(function, argument="An argument"), but I was wrong. What is it really? Well, function = decorator(argument="An argument")(function) of course! Now this makes sense, and makes decorators much easier to parse. This type of decorator is easier to accomplish with a class-style decorator, which I won't get into here. But how would we make a decorator with optional arguments? It's a little more tricky. There are two "equivalent" statements that can express these decorators:
# Without arguments:
function = decorator(function)

# With arguments:
function = decorator(argument="Argument")(function)
So sometimes, our decorator will receive the function it is decorating. Other times, it will receive a keyword argument, and then its return value will be called with the function that we want to decorate.
I chose to use just a function-style decorator, like what I've been demonstrating above. This is how I accomplished it:
def decorator(func=None, **options):
    if func != None:
        # We received the function on this call, so we can define
        # and return the inner function
        def inner(*args, **kwargs):
            if len(options) > 0:
                print "Decorated function with options:"
            else:
                print "Decorated function without options!"
    
            for k, v in options.items():
                print "\t{}: {}".format(k, v)
    
            func(*args, **kwargs)
   
        return inner
  
    else:
        # We didn't receive the function on this call, so the return value
        # of this call will receive it, and we're getting the options now.
        def partial_inner(func):
            return decorator(func, **options)   
        return partial_inner
    
@decorator
def function_a(x, y):
    print x + y
 
@decorator(foo="bar", baz=42)
def function_b(x, y):
    print x + y
 
function_a(2, 3)
#> Decorated function without options!
#> 5

function_b(2, 3)
#> Decorated function with options:
#>    foo: bar
#>    baz: 42
#> 5
You can see that I've accounted for the fact that if the decorator function is called without the function to be decorated, its return value will get the function instead. To handle this, I created another inner function, partial_inner(), to capture the function and re-call the decorator with the options. Note that similar to the inner() function, partial_inner() closes its outer scope in the decorator, so when it is called with the function to be decorated it still has the options passed the first time available.
This solution has a caveat: it only works if what you're passing to the decorator are keyword arguments. If you use regular positional arguments, the first one will be put into the func argument of the decorator, and you'll probably get the error that it isn't callable. If you want to pass an optional, callable, positional argument to a decorator, it gets even more convoluted. I find this keywords-only solution to be acceptable, and it works well for my purposes. I hope others find this useful as well.

2 comments:

  1. For some reason this doesn't work. Pasting your code in leaves function_b with the value None. I'm using Python 2.7, perhaps you are using Python 3?

    ReplyDelete
  2. Whoops! Of course it does, the `partial_inner` function wasn't returning a value :/ I've fixed it.

    ReplyDelete