Creating a Python decorator with optional arguments

On my current project I have created a small Python client library for talking to a remote REST API. The library is designed to provide an easy to use interface for other developers whilst shielding them from the complexity of user authentication and connection retry behaviour.

The REST commands that the library exposes are pluggable. Any callable can be registered with the library as a command, termed an “operation”.

For registering a callable as an operation, I wanted to use a decorator. As well as providing the registration behaviour, a decorator would have the benefit of clearly communicating the purpose of a callable it was decorating. I named the decorator @operation which would be used as follows:

@operation
def ingest_data():
    # Operation that ingests some data,
    # implemented as a function

The decorator could also be used to decorate class-based operations that have a __call__method:

@operation
class IngestData:
    def __call__(self):
        # Operation that ingests some data, 
        # implemented as a callable object

The @operation decorator would register a callable by simply adding it to a dictionary, using the callable’s lower case __name__ as the key. That name could then used when presenting a list of registered operations to a user, as part of the library’s help functionality. The decorator would just return the callable after registering it.

operations_registry = {}

def operation(callable_):
    if isinstance(callable_, type):
        op = callable_()
    else:
        op = callable_ 

    operations_registry[callable_.__name__.lower()] = op

    return callable_

But then I had a thought – what if somebody wanted to use an operation name different to the callable’s __name__? What if __name__ wasn’t appropriate? I decided that the @operationdecorator should support an optional name argument which, when supplied, would be used in preference to __name__.

@operation(name='trigger_ingestion')
def ingest_data():
    # Operation that specifies an explicit name

Supporting optional arguments

Decorators that accept arguments work in a subtly different way to decorators that do not. A decorator that takes no arguments receives the decorated function directly. A decorator that takes arguments must return a function which itself then receives the decorated function. i.e. a decorator that takes arguments must return a decorator.

So how to support both forms of @operation and @operation(name='some_name') elegantly in the same decorator implementation? I consulted the excellent Python Cookbook 3rd Edition, which has a very good solution using functools.partial.

import functools

operations_registry = {}

# Based on recipe 9.6 from the Python Cookbook 3rd Edition.

def operation(callable_=None, name=None):
    if callable_ is None:
       return functools.partial(operation, name=name)

    if name is not None:
        key = name
    else:
        key = callable_.__name__.lower()

    if isinstance(callable_, type):
        op = callable_()
    else:
        op = callable_

    operations_registry[key] = op

    return callable_

Notice that the above solution uses default values for both the 
name and callable_ arguments, allowing both of them to be optional.

When the name argument is supplied to the decorator, the decorator function will first be invoked with just the name (callable_ will be None) and it will immediately return a version of itself with the name argument fixed using functools.partial. That version will then be invoked with just the callable_ argument which it will register using the value of name that was fixed from the first invocation.

When the name argument is not supplied, the decorator will be only invoked once (callable_will be set, name will be None) and it will work exactly like before, registering the callable using its __name__ attribute.

Nice.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s