Overview
In the article Deep Dive Into Python Decorators, I introduced the concept of Python decorators, demonstrated many cool decorators, and explained how to use them.
In this tutorial I’ll show you how to write your own decorators. As you’ll see, writing your own decorators gives you a lot of control and enables many capabilities. Without decorators, those capabilities would require a lot of error-prone and repetitive boilerplate that clutters your code or completely external mechanisms like code generation.
A quick recap if you know nothing about decorators. A decorator is a callable (function, method, class or object with a call() method) that accepts a callable as input and returns a callable as output. Typically, the returned callable does something before and/or after calling the input callable. You apply the decorator by using the @
The Hello World Decorator
Let’s start with a ‘Hello world!’ decorator. This decorator will totally replace any decorated callable with a function that just prints ‘Hello World!’.
def hello_world(f):
def decorated(*args, **kwargs):
print 'Hello World!'
return decorated
That’s it. Let’s see it in action and then explain the different pieces and how it works. Suppose we have the following function that accepts two numbers and prints their product:
def multiply(x, y):
print x * y
If you invoke, you get what you expect:
multiply(6, 7) 42
Let’s decorate it with our hello_world decorator by annotating the multiply function with @hello_world.
@hello_world
def multiply(x, y):
print x * y
Now, when you call multiply with any arguments (including wrong data types or wrong number of arguments), the result is always ‘Hello World!’ printed.
multiply(6, 7)
Hello World!
multiply()
Hello World!
multiply('zzz')
Hello World!
OK. How does it work? The original multiply function was completely replaced by the nested decorated function inside the hello_world decorator. If we analyze the structure of the hello_world decorator then you’ll see that it accepts the input callable f (which is not used in this simple decorator), it defines a nested function called decorated that accepts any combination of arguments and keyword arguments (def decorated(*args, **kwargs)), and finally it returns the decorated function.
Writing Function and Method Decorators
There is no difference between writing a function and a method decorator. The decorator definition will be the same. The input callable will be either a regular function or a bound method.
Let’s verify that. Here is a decorator that just prints the input callable and type before invoking it. This is very typical for a decorator to perform some action and continue by invoking the original callable.
def print_callable(f):
def decorated(*args, **kwargs):
print f, type(f)
return f(*args, **kwargs)
return decorated
Note the last line that invokes the input callable in a generic way and returns the result. This decorator is non-intrusive in the sense that you can decorate any function or method in a working application, and the application will continue to work because the decorated function invokes the original and just has a little side effect before.
Let’s see it in action. I’ll decorate both our multiply function and a method.
@print_callable
def multiply(x, y):
print x * y
class A(object):
@print_callable
def foo(self):
print 'foo() here'
When we call the function and the method, the callable is printed and then they perform their original task:
multiply(6, 7) <function multiply at 0x103cb6398> <type 'function'> 42 A().foo() <function foo at 0x103cb6410> <type 'function'> foo() here
Decorators With Arguments
Decorators can take arguments too. This ability to configure the operation of a decorator is very powerful and allows you to use the same decorator in many contexts.
Suppose your code is way too fast, and your boss asks you to slow it down a little bit because you’re making the other team members look bad. Let’s write a decorator that measures how long a function is running, and if it runs in less than a certain number of seconds t, it will wait until t seconds expire and then return.
What is different now is that the decorator itself takes an argument t that determines the minimum runtime, and different functions can be decorated with different minimum runtimes. Also, you will notice that when introducing decorator arguments, two levels of nesting are required:
import time
def minimum_runtime(t):
def decorated(f):
def wrapper(*args, **kwargs):
start = time.time()
result = f(*args, **kwargs)
runtime = time.time() - start
if runtime < t:
time.sleep(t - runtime)
return result
return wrapper
return decorated
Let’s unpack it. The decorator itself—the function minimum_runtime takes an argument t, which represents the minimum runtime for the decorated callable. The input callable f was “pushed down” to the nested decorated function, and the input callable arguments were “pushed down” to yet another nested function wrapper.
The actual logic takes place inside the wrapper function. The start time is recorded, the original callable f is invoked with its arguments, and the result is stored. Then the runtime is checked, and if it’s less than the minimum t then it sleeps for the rest of the time and then returns.
To test it, I’ll create a couple of functions that call multiply and decorate them with different delays.
@minimum_runtime(1)
def slow_multiply(x, y):
multiply(x, y)
@minimum_runtime(3)
def slower_multiply(x, y):
multiply(x, y)
Now, I’ll call multiply directly as well as the slower functions and measure the time.
import time
funcs = [multiply, slow_multiply, slower_multiply]
for f in funcs:
start = time.time()
f(6, 7)
print f, time.time() - start
Here is the output:
42 <function multiply at 0x103cb6b90> 1.59740447998e-05 42 <function wrapper at 0x103d0bcf8> 1.00477004051 42 <function wrapper at 0x103cb6ed8> 3.00489807129
As you can see, the original multiply took almost no time, and the slower versions were indeed delayed according to the provided minimum runtime.
Another interesting fact is that the executed decorated function is the wrapper, which makes sense if you follow the definition of the decorated. But that could be a problem, especially if we’re dealing with stack decorators. The reason is that many decorators also inspect their input callable and check its name, signature and arguments. The following sections will explore this issue and provide advice for best practices.
Object Decorators
You can also use objects as decorators or return objects from your decorators. The only requirement is that they have a __call__() method, so they are callable. Here is an example for an object-based decorator that counts how many times its target function is called:
class Counter(object):
def __init__(self, f):
self.f = f
self.called = 0
def __call__(self, *args, **kwargs):
self.called += 1
return self.f(*args, **kwargs)
Here it is in action:
@Counter
def bbb():
print 'bbb'
bbb()
bbb
bbb()
bbb
bbb()
bbb
print bbb.called
3
Choosing Between Function-Based and Object-Based Decorators
This is mostly a question of personal preference. Nested functions and function closures provide all the state management that objects offer. Some people feel more at home with classes and objects.
In the next section, I’ll discuss well-behaved decorators, and object-based decorators take a little extra work to be well-behaved.
Well-Behaved Decorators
General-purpose decorators can often be stacked. For example:
@decorator_1
@decorator_2
def foo():
print 'foo() here'
When stacking decorators, the outer decorator (decorator_1 in this case) will receive the callable returned by the inner decorator (decorator_2). If decorator_1 depends in some way on the name, arguments or docstring of the original function and decorator_2 is implemented naively, then decorator_2 will see not see the correct information from the original function, but only the callable returned by decorator_2.
For example, here is a decorator that verifies its target function’s name is all lowercase:
def check_lowercase(f):
def decorated(*args, **kwargs):
assert f.func_name == f.func_name.lower()
f(*args, **kwargs)
return decorated
Let’s decorate a function with it:
@check_lowercase
def Foo():
print 'Foo() here'
Calling Foo() results in an assertion:
In [51]: Foo()
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
<ipython-input-51-bbcd91f35259> in <module>()
----> 1 Foo()
<ipython-input-49-a80988798919> in decorated(*args, **kwargs)
1 def check_lowercase(f):
2 def decorated(*args, **kwargs):
----> 3 assert f.func_name == f.func_name.lower()
4 return decorated
But if we stack the check_lowercase decorator over a decorator like hello_world that returns a nested function called ‘decorated’ the result is very different:
@check_lowercase
@hello_world
def Foo():
print 'Foo() here'
Foo()
Hello World!
The check_lowercase decorator didn’t raise an assertion because it didn’t see the function name ‘Foo’. This is a serious problem. The proper behavior for a decorator is to preserve as much of the attributes of the original function as possible.
Let’s see how it’s done. I’ll now create a shell decorator that simply calls its input callable, but preserves all the information from the input function: the function name, all its attributes (in case an inner decorator added some custom attributes), and its docstring.
def passthrough(f):
def decorated(*args, **kwargs):
f(*args, **kwargs)
decorated.__name__ = f.__name__
decorated.__name__ = f.__module__
decorated.__dict__ = f.__dict__
decorated.__doc__ = f.__doc__
return decorated
Now, decorators stacked on top of the passthrough decorator will work just as if they decorated the target function directly.
@check_lowercase
@passthrough
def Foo():
print 'Foo() here'
Using the @wraps Decorator
This functionality is so useful that the standard library has a special decorator in the functools module called ‘wraps’ to help write proper decorators that work well with other decorators. You simply decorate inside your decorator the returned function with @wraps(f). See how much more concise passthrough looks when using wraps:
from functools import wraps
def passthrough(f):
@wraps(f)
def decorated(*args, **kwargs):
f(*args, **kwargs)
return decorated
I highly recommend always using it unless your decorator is designed to modify some of these attributes.
Writing Class Decorators
Class decorators were introduced in Python 3.0. They operate on an entire class. A class decorator is invoked when a class is defined and before any instances are created. That allows the class decorator to modify pretty much every aspect of the class. Typically you’ll add or decorate multiple methods.
Let’s jump right in to a fancy example: suppose you have a class called ‘AwesomeClass’ with a bunch of public methods (methods whose name doesn’t start with an underscore like init) and you have a unittests-based test class called ‘AwesomeClassTest’. AwesomeClass is not just awesome, but also very critical, and you want to ensure that if someone adds a new method to AwesomeClass they also add a corresponding test method to AwesomeClassTest. Here is the AwesomeClass:
class AwesomeClass:
def awesome_1(self):
return 'awesome!'
def awesome_2(self):
return 'awesome! awesome!'
Here is the AwesomeClassTest:
from unittest import TestCase, main
class AwesomeClassTest(TestCase):
def test_awesome_1(self):
r = AwesomeClass().awesome_1()
self.assertEqual('awesome!', r)
def test_awesome_2(self):
r = AwesomeClass().awesome_2()
self.assertEqual('awesome! awesome!', r)
if __name__ == '__main__':
main()
Now, if someone adds an awesome_3 method with a bug, the tests will still pass because there is no test that calls awesome_3.
How can you ensure that there is always a test method for every public method? Well, you write a class decorator, of course. The @ensure_tests class decorator will decorate the AwesomeClassTest and will make sure every public method has a corresponding test method.
def ensure_tests(cls, target_class):
test_methods = [m for m in cls.__dict__ if m.startswith('test_')]
public_methods = [k for k, v in target_class.__dict__.items()
if callable(v) and not k.startswith('_')]
# Strip 'test_' prefix from test method names
test_methods = [m[5:] for m in test_methods]
if set(test_methods) != set(public_methods):
raise RuntimeError('Test / public methods mismatch!')
return cls
This looks pretty good, but there is one problem. Class decorators accept just one argument: the decorated class. The ensure_tests decorator needs two arguments: the class and the target class. I couldn’t find a way to have class decorators with arguments similar to function decorators. Have no fear. Python has the functools.partial function just for these cases.
@partial(ensure_tests, target_class=AwesomeClass)
class AwesomeClassTest(TestCase):
def test_awesome_1(self):
r = AwesomeClass().awesome_1()
self.assertEqual('awesome!', r)
def test_awesome_2(self):
r = AwesomeClass().awesome_2()
self.assertEqual('awesome! awesome!', r)
if __name__ == '__main__':
main()
Running the tests results in success because all the public methods, awesome_1 and awesome_2, have corresponding test methods, test_awesome_1 and test_awesome_2.
---------------------------------------------------------------------- Ran 2 tests in 0.000s OK
Let’s add a new method awesome_3 without a corresponding test and run the tests again.
class AwesomeClass:
def awesome_1(self):
return 'awesome!'
def awesome_2(self):
return 'awesome! awesome!'
def awesome_3(self):
return 'awesome! awesome! awesome!'
Running the tests again results in the following output:
python3 a.py
Traceback (most recent call last):
File "a.py", line 25, in <module>
class AwesomeClassTest(TestCase):
File "a.py", line 21, in ensure_tests
raise RuntimeError('Test / public methods mismatch!')
RuntimeError: Test / public methods mismatch!
The class decorator detected the mismatch and notified you loud and clear.
Conclusion
Writing Python decorators is a lot of fun and lets you encapsulate tons of functionality in a reusable way. To take full advantage of decorators and combine them in interesting ways, you need to be aware of best practices and idioms. Class decorators in Python 3 add a whole new dimension by customizing the behavior of complete classes.
Comments