Write Your Own Python Decorators


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 @ syntax. Plenty of examples coming soon...

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!’.

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:

If you invoke, you get what you expect:

Let’s decorate it with our hello_world decorator by annotating the multiply function with @hello_world.

Now, when you call multiply with any arguments (including wrong data types or wrong number of arguments), the result is always ‘Hello World!’ printed.

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.

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.

When we call the function and the method, the callable is printed and then they perform their original task:

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:

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.

Now, I’ll call multiply directly as well as the slower functions and measure the time.

Here is the output:

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:

Here it is in action:

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:

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:

Let’s decorate a function with it:

Calling Foo() results in an assertion:

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:

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.

Now, decorators stacked on top of the passthrough decorator will work just as if they decorated the target function directly.

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:

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:

Here is the AwesomeClassTest:

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.

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.

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.

Let’s add a new method awesome_3 without a corresponding test and run the tests again.

Running the tests again results in the following output:

The class decorator detected the mismatch and notified you loud and clear.


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.



Related Articles