Decorators with Optional Arguments in Python

It is common that features dubbed ‘syntactic sugar’ are often fostering novel approaches to programming problems. Python’s decorators are no different here, and this was a topic I touched upon before. Today I’d like to discuss few quirks which are, unfortunately, adding to their complexity in a way that often doesn’t feel necessary.

Let’s start with something easy. Pretend that we have a simple decorator named @trace, which logs every call of the function it is applied to:

An implementation of such decorator is relatively trivial, as it wraps the decorated function directly. One of the possible variants can be seen below:

That’s pretty cool for starters, but let’s say we want some calls to stand out in the logging output. Perhaps there are functions that we are more interested in than the rest. In other words, we’d like to adjust the priority of log messages that are generated by @trace:

This seemingly small change is actually mandating massive conceptual leap in what our decorator really does. It becomes apparent when we de-sugar the @decorator syntax and look at the plumbing underneath:

Introduction of parameters requires adding a new level of indirection, because it’s the return value of trace(level=logging.INFO) that does the actual decorating (i.e. transforming given function into another). This might not be obvious at first glance and admittedly, a notion of function that returns a function which takes some other function in order to output a final function might be – ahem – slightly confusing ;-)

But wait! There is just one more thing… When we added the level argument, we not necessarily wanted to lose the ability to invoke @trace without it. Yes, it is still possible – but the syntax is rather awkward:

That’s expected – trace only returns the actual decorator now – but at least slightly annoying. Can we get our pretty syntax back while maintaining the added flexibility of specifying custom arguments? Better yet: can we make @trace, @trace() and @trace(level) all work at the same time?…

Looks like tough call, but fortunately the answer is positive. Before we delve into details, though, let’s step back and try to somewhat improve the way we are writing our decorators.

A class is generally better

As we probably know very well, the quacky nature of Python allows not only functions to be callable. Other language constructs which can also behave that way include all classes, as well as certain kind of objects – namely those implementing the __call__ method. Stitching those two facts together, we can see that classes having __call__ method can be considered “doubly callable”:

This is strikingly similar to de-sugared version of our parametrized @trace decorator. And indeed, it is relatively straightforward to implement as a class:

For experienced Python developer, the difference might be small, but personally I find it much easier to wrap my head around such class rather than triply-nested function. Also, note that it serves us nicely by clearly separating the key two steps of parametrized decoration: creating the decorator object in __init__ (with supplied or default argument) and applying it to a function in __call__. This will come handy in our quest for reclaiming the neat syntax for decorators with optional parameters.

Are you a function, maybe?

To properly support both @trace and @trace(...) form of our decorator, we must properly recognize those two situations. This boils down to determining whether the decorator is directly applied, or is just getting its __init__ parameters. Both cases will provoke a call, but with different arguments – and it’s a difference we can generally detect:

It’s rather simple, too. If the decorator is invoked directly, it receives a function as its sole positional parameter. This function shall be decorated, so we are instantiating our class (with default arguments) and “calling” its object on this single parameter that we have received.
On the other hand, getting something else indicates that the decorator is not being called to wrap a function, but is just having some optional arguments passed. In this case we should simply relay them to the class and return an object which is created as a result. This situation corresponds to the parametrized invocation of our decorator, such as @trace(level=logging.INFO).

Note that if we are writing universal decorator that works for any callable, we should obviously replace inspect.isfunction with callable for our heuristics to work correctly.

Getting more general

This pattern of dispatching (described above using the @trace example) is general enough to be used by pretty much any parametrized decorator. It doesn’t even have to be implemented as a class; it just needs to be “doubly callable” and allow for its arguments to be optional. We don’t even need to touch the decorator itself, as the pattern can be encapsulated in a function of its own:

Applying it to the TraceDecorator is trivial:

It actually goes even further, as the optional_args function is itself capable of working as decorator! Provided we use at least Python 2.6, we can put it directly in front of our TraceDecorator definition:

In other words, we are decorating our decorators, so we can more easily decorate other things with them. This just screams meta-awesomeness :)

Wrapping up

There is one small limitation of this technique for detecting parametrized and parameterless invocations of decorators. It’s the case when decorator can accept a single positional argument which is a function – other than the function being decorated, that is:

It seems to be rather contrived, though. And even if this issue arises in practice, solution would be very simple. The problematic argument shall be simply made keyword rather than positional:

This causes it to land inside kwargs dictionary and prevents from being incorrectly recognized as decorated function.

Be Sociable, Share!

• 
• 

Be Sociable, Share!

• 
• 

• Tweet
•