abidibo.net

Decorators with Python

decorator patterns python

Decorators in python can be quite easy to use and understand, but flexible decorators using the new syntax introduced by python 2.4 can be a bit more complex.

Here I'll try to explain (especially to myself) the world of decorators in python, starting from simple cases and ending with complex ones. Ideas and notions taken from the Pro Django book by Marty Alchin.

Decorators without extra arguments

Let's create a simple cache decorator you can use to store complex operations:

>>> def cache(func):
...   cache_dict = {}
...   def wrapper(*args, **kwargs):
...     key = '%s-%s' % (repr(args), repr(kwargs))
...     if key not in cache_dict:
...       cache_dict[key] = func(*args, **kwargs)
...       print 'cached'
...     return cache_dict[key]
...   return wrapper

A simple decorator function takes as an argument the function that needs to be decorated, and returns another function which simply adds some functionality to the original one.

In this case we define a cache dict which will store the cached values. It is defined outside the returned function so that it is accessible from the wrapper function which retains it,  this is a closure.
The key is unique for different set of arguments, then the wrapper function checks if the function to be decorated was already called with the given arguments, in such case it returns the stored value, otherwise it executes the function and stores the value.

Use it the old way

>>> def sum(a, b):
...   return a + b
... 
>>> sum = cache(sum)
>>> sum(1,3)
cached
4
>>> sum(1,3)
4

A function is defined, then is passed as an argument to the decorator function which returns the wrapper function, which is the original function decorated with extra functionalities.

As you can see the first time the function is called the result is cached, the second time the stored value is returned.

New syntax

>>> @cache
... def sum(a, b):
...   return a + b
...
>>> sum(1,5)
cached
6
>>> sum(1,5)
6

The new syntax is cooler without doubt, but it has a drawback. In fact what about if you need to pass an extra argument to the decorator function?

It will still be simple in the old way since the decorator function is called directly:

sum = cache(sum, arg = 'my extra argument')

but with the new syntax python will call the decorator function passing only the decorated function as the first positional argument. Fortunately there are some techniques we can use to overcome such problem. Let's see them.

Decorators with extra arguments

Consider this code:

>>> def cache(msg='cached'):
...   def decorator(func):
...     cache_dict = {}
...     def wrapper(*args, **kwargs):
...       key = '%s-%s' % (repr(args), repr(kwargs))
...       if key not in cache_dict:
...         cache_dict[key] = func(*args, **kwargs)
...         print msg
...       return cache_dict[key]
...     return wrapper
...   return decorator
... 
>>> @cache(msg='my custom msg')
... def sum(a, b):
...   return a + b
... 
>>> sum(2,3)
my custom msg
5
>>> sum(2,3)
5
>>> @cache()
... def sum(a, b):
...   return a + b
... 
>>> sum(2,3)
cached
5
>>> sum(2,3)
5
>>> @cache
... def sum(a, b):
...   return a + b
... 
>>> sum(2, 3)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: decorator() takes exactly 1 argument (2 given)

The trick here is to define the decorator inside another function that will accept the extra argument. Then such function returns the real decorator function which can use the extra argument because it is in its scope.

As a drawback we can't call the cache decorator without parentheses, If we do not need the extra argument we still have to call it with them, we still need to execute the function. If we don't an error occurs, since python will call the decorator function passing the arguments provided to the decorated function (sum), in my example 2, 3.

As often happens, we can do better and allow the use of the decorator without parentheses, we just need to change the code a bit.

Decorators with or without extra arguments

The trick is to check how the decorator function is called, by checking if the first positional argument we define, the func argument, is provided or not. When using the decorator passing extra arguments it will not be provided, the opposite occurs when used without parentheses.

>>> def cache(func=None, msg='cached'):
...   def decorate(func):
...     cache_dict = {}
...     def wrapper(*args, **kwargs):
...       key = '%s-%s' % (repr(args), repr(kwargs))
...       if key not in cache_dict:
...         cache_dict[key] = func(*args, **kwargs)
...         print msg
...       return cache_dict[key]
...     return wrapper
...   if func is None:
...     # called with arguments
...     def decorator(func):
...       return decorate(func)
...     return decorator
...   # called without arguments 
...   return decorate(func)
... 
>>> @cache
... def sum(a, b):
...   return a + b
... 
>>> sum(2,3)
cached
5
>>> sum(2,3)
5
>>> @cache(msg='my custom msg')
... def sum(a, b):
...   return a + b
... 
>>> sum(2, 3)
my custom msg
5
>>> sum(2, 3)
5

If cache is called with arguments a decorator function is returned, which can receive a function argument, and returns the decorated function (wrapper).

If cache is called without arguments, and then the function to be decorated is provided, returns directly the decorated function (wrapper).

Notice that the extra arguments must be defined as keyword arguments!

If you're wondering why you need to wrap the wrapper function inside the decorate function, the answer is that otherwise returning directly the wrapper function, the original func would be undefined inside it.

Every time we want to define a decorator which can accept or not extra arguments we should code something like that, which can be annoying, so the final step is to write a decorator which decorates a decorator giving it the property of being callable with or without arguments, let's see how.

The trick is to create a decorator which decorates the real decorator (which implements the wrapper function in my example) and returns a function like the one seen in the previuos example.

>>> def optional_arg_decorator(real_decorator):
...   def decorator(func=None, **kwargs):
...     def decorate(func):
...       def wrapper(*a, **kw):
...         return real_decorator(func, a, kw, **kwargs)
...       return wrapper
...     if func is None:
...       def decorator(func):
...         return decorate(func)
...       return decorator
...     return decorate(func)
...   return decorator
... 
>>> cached_dict = {}
>>> @optional_arg_decorator
... def cache(func, args, kwargs, msg='cached'):
...   key = '%s-%s' % (repr(args), repr(kwargs))
...   if key not in cached_dict:
...     cached_dict[key] = func(*args, **kwargs)
...     print msg
...   return cached_dict[key]
... 
>>> @cache
... def sum(a, b):
...   return a + b
... 
>>> sum(2, 3)
cached
5
>>> sum(2, 3)
5
>>> @cache(msg='my custom msg')
... def sum(a, b):
...   return a + b
... 
>>> sum(2, 3)
5
>>> sum(2, 5)
my custom msg
7
>>> sum(2, 5)
7

The difference is that the real decorator function just executes the inner code of the previous examples' wrapper function, and returns the result. This happens because it is itself decorated, and its decorated form returns exactly that wrapper function which now call the real_decorator to run specific code.

As you can see the cache_dict obj now was defined outside the decorator, since it can't stay inside it (it would be reset every time). So it would be necessary to create another closure in a real world example, but here I just want to clarify how decorators work. As a side effect of such thing, you can see that a second decoration doesn't reset the cache dictionary, so the first call to sum(2, 3) after the second decoration still is cached.

Well I found many difficulties to explain these concepts in a language which is not my native one, but I hope someone can find it useful.

As usual comments, improvements, opinions are appreciated.

Subscribe to abidibo.net!

If you want to stay up to date with new contents published on this blog, then just enter your email address, and you will receive blog updates! You can set you preferences and decide to receive emails only when articles are posted regarding a precise topic.

I promise, you'll never receive spam or advertising of any kind from this subscription, just content updates.

Subscribe to this blog

Comments are welcome!

blog comments powered by Disqus

Your Smartwatch Loves Tasker!

Your Smartwatch Loves Tasker!

Now available for purchase!

Featured