Updated: 2021-01-28 added corner cases for class-based decorators

This article is an overview of Python decorators trying to cover many aspects starting from the basic steps and usages and moving into complex class-based decorators for asyncronious programming. Decorator is a simple yet powerful syntatic sugar in Python language that can make a life of a developer easier and keep codebase clean.

The official Python 3 documentation provides a very short description of decorators:

A function returning another function, usually applied as a function transformation using the @wrapper syntax. Common examples for decorators are classmethod() and staticmethod().

The decorator syntax is merely syntactic sugar, the following two function definitions are semantically equivalent:

def f(...):
    ...
f = staticmethod(f)

@staticmethod
def f(...):
    ...

The same concept exists for classes, but is less commonly used there. See the documentation for function definitions and class definitions for more about decorators.

Though this description is valid it provides very little information about the true abilities of the subject. And that’s what we’ll try to fix.

Those that prefer reading source code to long texts may check out fqn-decorators Python library and its source code on GitHub to understand how class-based sync and async parametrized decorators can be created and extended.

Index

Creating the simplest decorator

Decorators are wrappers around functions and classes providing additional functionality around the wrapped objects. To understand how decorators work first we need to look into closures. Closure is a way to create functions with a specific scope. For example:

def wait_generator(sec):
    def inner():
        print(f'waiting for {sec} seconds...')
        time.sleep(sec)

    return inner

Note, that inner() function uses sec argument, which is not passed directly into it, but available for it in scope of wait_generator() function. This provides a template that allows to create different wait functions:

wait_2_sec = wait_generator(2)
wait_3_sec = wait_generator(3)

wait_2_sec()
>>> waiting for 2 seconds...
wait_3_sec()
>>> waiting for 3 seconds...

The same way we pass sec as argument to a wait_generator we can pass another function. Let’s create a function that will take another method as agrument and print some nice message before and after its execution.

def execute_function(func):
    def inner():
        print(f'About to run "{func.__name__}()"')
        func()
        print(f'"{func.__name__}()" run complete')

    return inner

And we’ll take a simple func that prints out “Hello”:

def hello():
    print('Hello')

Now we can combine them together:

wrapped_hello = execute_function(hello)

wrapped_hello()
>>> About to run "hello()"
>>> Hello
>>> "hello()" run complete

To make this closure be a real decorator we need to reassign the result of the execute_function to the same func we passed as argument:

hello()  # before wrapper function
>>> Hello

hello = execute_function(hello)

hello()  # after wrapper function
>>> About to run "hello()"
>>> Hello
>>> "hello()" run complete

That’s it! We created our first decorator: execute_function, because this:

def hello():
    print('Hello')

hello = execute_function(hello)

can be rewritten as:

@execute_function
def hello():
    print('Hello')

Getting more real

Handle function arguments and return value

We can make our example more real-life to support func to have arguments and also make the wrapper return the result of the func. We’ll create a function that returns the mathematical sum of all its arguments. Note, that it’s specifically made non-optimised and slow for demonstration purpose only (also because I didn’t think of an interesting one):

def sum_numbers(*args):
    return sum(args)

But our execute_function does not know in advance what func will be passed and how many arguments it can have, so we make it as generic as possible:

def execute_function(func):
    def inner(*args, **kwargs):
        print(f'About to run "{func.__name__}()"')
        result = func(*args, **kwargs)
        print(f'"{func.__name__}()" run complete')
        return result

    return inner

Note, that execute_function signature doesn’t change: it still takes only func as argument, but inner now supports to have arguments that will be passed to func:

sum_numbers = execute_function(sum_numbers)

result = sum_numbers(5, 4)
>>> About to run "sum_numbers()"
>>> "sum_numbers()" run complete
print(result)
>>> 9

Decorator to measure function execution time

Now that we know how decorators work let’s modify execute_function to measure inner function execution time instead of just priting lines of text.

import time

def measure_time(func):
    def inner(*args, **kwargs):
        start_t = time.time()
        result = func(*args, **kwargs)
        end_t = time.time()
        print(f'Execution time: {end_t - start_t}s')
        return result

    return inner

And we can apply it to a function to record its run time:

@measure_time
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)  # 15
>>> Execution time: 2.86102294921875e-06s
sum_numbers(*range(int(1E+06)))  # 499999500000
>>> Execution time: 0.022812366485595703s

Parametrized decorator

For now our decotrator always prefixes the output string with “Execution time” phrase, but what if we want to make it configurable? E.g. consider this example:

@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)
>>> Time spent: 2.87491942710492e-06s

In order to achieve this we need to remember that decorator in nothing more than a function itself and it can handle parameters as any other function. First let’s remind us how decorator without any arguments looks like. We’ll write it first using decorator syntax and then its corresponding function syntax below:

@measure_time
def func(*args, **kwargs):
   ...

func = measure_time(func)

Then decorator with arguments should look like:

@measure_time(text='Time spent')
def func(*args, **kwargs):
   ...

func = measure_time(text='Time spent')(func)

This is a very important moment here: now measure_time should accept argument(s) and instead of being a decorator should return a decorator method. So a code for measure_time will look like this:

import time

def measure_time(text=None):
    text = text if text is not None else 'Execution time'
    def internal(func):
        def wrapper(*args, **kwargs):
            start_t = time.time()
            result = func(*args, **kwargs)
            end_t = time.time()
            print(f'{text}: {end_t - start_t}s')
            return result
        return wrapper
    return internal

@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)
>>> Time spent: 1.9073486328125e-06s

# As we made `text` optional it's possible to run decorator without arguments.
# Note that parentheses are still required.
@measure_time()
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)
>>> Execution time: 2.86102294921875e-06s

Advanced usage

Decorator with optional arguments

So we created a decorator that accepts an argument, but now we have a problem: it’s not possible to use decorator without parentheses because only the call to our measure_time function returns a decorator. Would be nice to be able to support both behaviours so we don’t need extra pair of parentheses when we don’t use any arguments:

@measure_time
def sum_numbers(sec):
    ...

@measure_time(text='Time spent')
def multiply_numbers(sec):
    ...

To distinguish between these two calls we can check what the first decorator argument is as in the first case it’ll be a function:

func = measure_time(func)

and in the second example - a text parameter:

func = measure_time(text='Time spent')(func)

Then the code for measure_time will be:

import time

def measure_time(*decorator_args, **decorator_kwargs):
    def internal(func):
        def wrapper(*args, **kwargs):
            start_t = time.time()
            result = func(*args, **kwargs)
            end_t = time.time()
            print(f'{text}: {end_t - start_t}s')
            return result
        return wrapper

    text = 'Execution time'
    if len(decorator_args) == 1 and callable(decorator_args[0]):
        return internal(decorator_args[0])

    text = decorator_kwargs.get('text', text)
    return internal

And now it’s possible to use both syntaxes:

@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)
>>> Time spent: 1.9073486328125e-06s

@measure_time
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6)
>>> Execution time: 2.86102294921875e-06s

Class-based decorator

Decorator is a syntactic sugar not only for functions but also for classes. In our last examples we had a measure_time function that returns internal that returns wrapper - not a very straight-forward code. Plus note that we are dealing with a very simple example to only measure function execution time. Real-life decorators are much complex and can make code hard to comprehend.

Class-based decorator can help make our code a bit cleaner and easier to read in this case. Its another feature is that being a class object in can store additional attributes between the function calls.

To understand how to create such class let’s look again at decorator shorcut syntax:

sum_numbers = measure_time(sum_numbers)
sum_numbers(5, 6, 7)

Imagine measure_time to be not a funciton, but a class. Then on the first line we create a class object, passing sum_numbers into class __init__ method. And on the second line we call this object with arguments.

sum_numbers = measure_time(sum_numbers)
              |                       |
              +-------__init__--------+

sum_numbers(5, 6, 7)
|                  |
+-----__call__-----+

So it means that a class that can be used as decorator should define two methods: __init__ and __call__:

class measure_time:
    # The decorated function is passed to class __init__ method
    def __init__(self, func):
        ...

    # __call__ is called when decorated function is executed
    def __call__(self, *args, **kwargs):
        ...

Let’s rewrtie out function-based decorator without any arguments into a class-based:

import time

def measure_time(func):
    def inner(*args, **kwargs):
        start_t = time.time()
        result = func(*args, **kwargs)
        end_t = time.time()
        print(f'Execution time: {end_t - start_t}s')
        return result
    return inner

class measure_time:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        start_t = time.time()
        result = self.func(*args, **kwargs)
        end_t = time.time()
        print(f'Execution time: {end_t - start_t}s')
        return result

Those two implemetations are identical and provide the same functionality:

@measure_time
def sum_numbers(*args):
    return sum(args)

sum_numbers(4, 5, 6, 7)
>>> Execution time: 2.86102294921875e-06s

Class-based decorator attributes

We already mentioned that class-based decorator object can store and use attributes between the function calls. It can be both a power and also a weakness (or a bug) if used incorrectly.

Let’s take another classic example of a decorator that should count how many times the method was called. It can’t be easily achieved with a function-based one, but that’s what a class-based is for:

class counter:
    def __init__(self, func):
        self.func = func
        # __init__ is called only once when decorator is applied to the method,
        # so we set the counter to 0 here
        self.call_count = 0

    def __call__(self, *args, **kwargs):
        self.call_count += 1  # each method call will increment the counter
        return self.func(*args, **kwargs)

and we take any function that we want to track and see how it goes:

@counter
def dummy_method():
    pass

print(dummy_method.__dict__)
# see how we can access decorator attributes now directly
>>> {'func': <function __main__.dummy_method()>, 'call_count': 0}

dummy_method()
print(dummy_method.call_count)
>>> 1

dummy_method()
print(dummy_method.call_count)
>>> 2

We will describe and deal with the drawbacks of this approach when we discuss class-based decorator with arguments and async methods.

Class-based decorator with arguments

Let’s rewrite our decorator with optional arguments into a class-based. First we’ll remind ourselves how a decorator with arguments looks like not in a short-cut form:

@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

Is the same as:

def sum_numbers(*args):
    return sum(args)

sum_numbers = measure_time(text='Time spent')(sum_numbers)

When we were dealing with functions measure_time was a function that produced a function that itself returned a function (not very pretty). Now that we switched to class-based logic, measure_time(text='Time spent') should be a class __init__, and (sum_numbers) call will be handled in __call__ logic being a usual simple function-based decorator without any arguments.

sum_numbers = measure_time(text='Time spent')(sum_numbers)
              |                             ||           |
              +---------__init__------------+|           |
                                             +--__call__-+

So the decorator class will look like this:

import time

class measure_time:
    def __init__(self, text=None):
        self.text = text or 'Execution time'

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            start_t = time.time()
            result = func(*args, **kwargs)
            end_t = time.time()
            print(f'{self.text}: {end_t - start_t}s')
            return result
        return wrapper

Now this class-based decorator can be used in a same way as a function-based:

@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

sum_numbers(1, 2, 3)
>>> Time spent: 1.9073486328125e-06s

or using the default output:

# Note that parentheses are required.
@measure_time()
def sum_numbers(*args):
    return sum(args)

sum_numbers(1, 2, 3)
>>> Execution time: 2.1457672119140625e-06s

Now let’s get back to our decorator that counts the method calls and make is also parametrized. Let’s make it print the passed text with current call clount, so we’ll have something like this:

dummy_method()
>>> "Method was called 12 times"
dummy_method()
>>> "Method was called 13 times"

Obviously we just add call_count attribute to __init__ as before, and printing and incrementing should go to the inner function call:

class counter:
    def __init__(self, text=None):
        self.text = text or 'Method was called {} times'
        self.call_count = 0

    def __call__(self, func):
        def wrapper(*args, **kwargs):
            self.call_count += 1
            print(self.text.format(self.call_count))
            return func(*args, **kwargs)
        return wrapper

Note that it is crucial to put the counter logic inside wrapper method because that’s the method that is called every time we call the decorated method, not the __call__ like in the case of decorator without arguments:

@counter(text="call count = {}")
def dummy_method(a):
    return a * 10

you may remember that is the same as

def dummy_method(a):
    return a * 10

dummy_method = counter(text="call count = {}")(dummy_method)
               |                             ||            |
               +---------__init__------------+|            |
                                              +--__call__--+

So both __init__ and __call__ are called only once when we apply the decorator. And when we actually call the method we are calling wrapper:

dummy_method(10)  # wrapper(10)
>>> call count = 1
>>> 100
dummy_method(20)  # wrapper(20)
>>> call count = 2
>>> 200

So if we make a mistake and try to increment the counter only in the __call__ like this:

    def __call__(self, func):
        self.call_count += 1
        def wrapper(*args, **kwargs):
            print(self.text.format(self.call_count))
            return func(*args, **kwargs)
        return wrapper

then our decorator with always print count 1:

dummy_method(10)
>>> call count = 1
>>> 100
dummy_method(20)
>>> call count = 1
>>> 200
dummy_method(30)
>>> call count = 1
>>> 300

We will discuss more shoot-in-your-foot possibilites when talking about async methods and decorators.

Class-based decorator with optional arguments

Now we ended up in a same not-very-pretty situation when using a decorator without any arguments we still need to use parentheses.

@measure_time()
def sum_numbers(*args):
    return sum(args)

Would be nice to rewrite our class-based decorator same way we did a function-based to support both usages:

@measure_time
def sum_numbers(*args):
    return sum(args)


@measure_time(text='Time spent')
def sum_numbers(*args):
    return sum(args)

To get the idea what we need to change let’s rewrite those decorator without any syntatic sugar:

# Let's call this CaseA
sum_numbers = measure_time(sum_numbers)
sum_numbers(4, 5, 6, 7)

and

# And this - CaseB
sum_numbers = measure_time(text='Time spent')(sum_numbers)
sum_numbers(4, 5, 6, 7)

It means that decorator __init__ should be able to handle both measure_time(sum_numbers) and measure_time(text='Time spent'). Then in CaseB the handling of measure_time(text='Time spent')(sum_numbers) should redirect the call back to CaseA __init__. This may sound a bit confusing, but the code below will show how it behaves:

import time

class measure_time:
    def __init__(self, func=None, text=None):
        # In CaseA func in not None and in CaseB - None
        self.func = func
        # In CaseA text is not present (the default one will be used)
        self.text = text or 'Execution time'

    def __call__(self, *args, **kwargs):
        if not self.func:
            # That will be CaseB:
            # the first and the only argument is a function to use the decorator
            # for; and we already saved a text to use in __init__.
            # So we emulate CaseA here but also passing text:
            return self.__class__(args[0], text=self.text)

        start_t = time.time()
        result = self.func(*args, **kwargs)
        end_t = time.time()
        print(f'{self.text}: {end_t - start_t}s')
        return result

This way it can handle both usages:

measure_time(text='Time spent')(sum_numbers)(1, 2, 3)
>>> Time spent: 1.9073486328125e-06s

# Note that empty parentheses are not required anymore
measure_time(sum_numbers)(1, 2, 3)
>>> Execution time: 1.6689300537109375e-06s

Decorator for a class

As you already guessed a decorator can be applied not only to functions but to classes. Same logic behind it applies. E.g. let’s create a decorator for a class that applies our measure_time functionality to all class methods:

@measure_time
class A:
    def x(self):
        return 10
    def y(self):
        return 20

a = A()
a.x()
>>> Execution time: 1.6689300537109375e-06s
a.y()
>>> Execution time: 1.6689300537109375e-06s

As before the decorator is a syntatic sugar:

class A:
    ...

A = measure_time(A)

That means that measure_time accepts class as an argument and should also return class modifying all its method in the process.

import time

def measure_time(cls):
    def method_decorator(func):
        def inner(*args, **kwargs):
            start_t = time.time()
            result = func(*args, **kwargs)
            end_t = time.time()
            print(f'Execution time: {end_t - start_t}s')
            return result
        return inner

    # Find all non-private class methods and apply method_decorator to them
    for func in filter(
        lambda func: not func.startswith('__') and callable(getattr(cls, func)),
        dir(cls)
    ):
        setattr(cls, func, method_decorator(getattr(cls, func)))
    return cls

Let’s put it to a test:

@measure_time
class A:
    x = 10
    def y(self):
        print(20)

a = A()
print(a.x)
>>> 10
a.y()
>>> 20
>>> Execution time: 9.5367431640625e-07s

Asynchronous decorators

Now we dive into Python asyncio library. We want to be able to apply our decorator to async funcitons like this:

@measure_time
async def sum_numbers(*args):
    return sum(args)

await sum_numbers(3, 4, 5)
>>> Execution time: 1.9489349310389372e-06s

We start again with function-based decorator without arguments. We need to be aware what to await and which methods will become async in decorator code now. Let’s write our async decorator without syntatic sugar:

await measure_time(sum_numbers)(3, 4, 5)

Now we can see that measure_time(sum_numbers) should not be async, but should return awaitable function.

import time

def measure_time(func):
    async def wrapper(*args, **kwargs):
        start_t = time.time()
        result = await func(*args, **kwargs)
        end_t = time.time()
        print(f'Execution time: {end_t - start_t}s')
        return result
    return wrapper

@measure_time
async def sum_numbers(*args):
    return sum(args)

await sum_numbers(3, 4, 5)
>>> Execution time: 1.9073486328125e-06s

Class-based async decorator with optional arguments

As you saw in a previous example async decorators don’t differ much from their sync versions. So we’ll discuss only the most complex case here. This example is similar to a corresponding synchronous one. We want to have a class-based async decorator that supports both usages:

@measure_time
async def sum_numbers(*args):
    return sum(args)

@measure_time(text='Time spent')
async def sum_numbers(*args):
    return sum(args)

Here again the first decorator call (__init__) should be sync and only __call__ - return an awaitable:

import time

class measure_time:
    def __init__(self, func=None, text=None):
        self.func = func
        self.text = text or 'Execution time'

    def __call__(self, *args, **kwargs):
        if not self.func:
            return self.__class__(args[0], text=self.text)

        async def wrapper(*args, **kwargs):
            start_t = time.time()
            result = await self.func(*args, **kwargs)
            end_t = time.time()
            print(f'{self.text}: {end_t - start_t}s')
            return result
        return wrapper(*args, **kwargs)

Now it can handle

@measure_time
async def sum_numbers(*args):
    return sum(args)

await sum_numbers(1, 2, 3)
>>> Execution time: 2.1457672119140625e-06s

and

@measure_time(text='Time spent')
async def sum_numbers(*args):
    return sum(args)

await sum_numbers(4, 5, 6)
>>> Time spent: 3.0994415283203125e-06s

Class-based async decorator with attributes and optional arguments

Now we’ve reached the most complex part of this article where we need all the knowledge above to handle this case. First let’s see what issue can happen with the async decorator. Let’s make our previous example a bit more real-life and add a couple of method to the decorator class before() and after() to be able to specify some generic logic to call before and after the actual function call and make them do the work related to time measurement:

import time

class measure_time:
    def __init__(self, func=None, text=None):
        self.func = func
        self.text = text or 'Execution time'

    def before(self):
        self.start_t = time.time()

    def after(self):
        end_t = time.time()
        print(f'{self.text}: {end_t - self.start_t}s')

    def __call__(self, *args, **kwargs):
        if not self.func:
            return self.__class__(args[0], text=self.text)

        async def wrapper(*args, **kwargs):
            self.before()
            result = await self.func(*args, **kwargs)
            self.after()
            return result
        return wrapper(*args, **kwargs)

So at the first sight nothing changed except we save the start time in a class attribute self.start_t to later use it in after() method. And if we use this decorator it will work just as before.

@measure_time
async def sleep(secs):
    await asyncio.sleep(sec)

await sleep(3)
>>> Execution time: 3.0048561096191406s

…or will it? Remember that we are now in async territory and same method can be executed in parallel in the same async loop. So what will happen when we run it several times in parallel with slight delay?

@measure_time
async def sleep_5_sec():
    await asyncio.sleep(5)  # we expect @measure_time to count 5 sec here

async def sleep_7_sec():
    await asyncio.sleep(2)
    await sleep_5_sec()  # @measure_time should also count 5 sec here

And then we execute them together at the same time:

await asyncio.gather(sleep_5_sec(), sleep_7_sec(), return_exceptions=True)
>>> Execution time: 3.00191068649292s  # What?! It should be 5 sec, not 3!
>>> Execution time: 5.004676818847656s

So what happened here? Why one counter appeared to be broken? Remember that we use class-based decorator object and this object is used for all calls and so are its attribures. Now we added this before() method:

def before(self):
    self.start_t = time.time()

def __call__(self, *args, **kwargs):
    ...

    async def wrapper(*args, **kwargs):
        self.before()
        result = await self.func(*args, **kwargs)
        ...

and as you know wrapper() is called upon each decorated method call, so each call will reset self.start_t to current time. sleep_7_sec() waited for 2 seconds while the first sleep_5_sec() was already running and reset the self.start_t to the current time which resulted in the first call to measure time wrong as 5 - 2 = 3 seconds.

So how can we fix this issue? One of the solutions would be to re-create a decorator instance for each invocation so that same instance is not shared between several parall runs anymore. We can add a flag to the decorator class, e.g. _initialized set to False by default and create a new class instance during the __call__ if decorator is not _initialized yet, setting this flag to True:

import time

class measure_time:
    def __init__(self, func=None, _initialized=False, text=None):
        self.func = func
        self.text = text or 'Execution time'
        self._initialized = _initialized

    def before(self):
        self.start_t = time.time()

    def after(self):
        end_t = time.time()
        print(f'{self.text}: {end_t - self.start_t}s')

    def __call__(self, *args, **kwargs):
        if not self._initialized:
            if not self.func:
                # Decorator was initialized with arguments
                return self.__class__(args[0], text=self.text)
            return self.__class__(self.func, _initialized=True, text=self.text)(*args, **kwargs)

        async def wrapper(*args, **kwargs):
            self.before()
            result = await self.func(*args, **kwargs)
            self.after()
            return result
        return wrapper(*args, **kwargs)

And let’s see how it works now:

@measure_time
async def sleep_5_sec():
    await asyncio.sleep(5)  # we expect @measure_time to count 5 sec here

async def sleep_7_sec():
    await asyncio.sleep(2)
    await sleep_5_sec()  # @measure_time should also count 5 sec here

await asyncio.gather(sleep_5_sec(), sleep_7_sec(), return_exceptions=True)
>>> Execution time: 5.002807140350342s
>>> Execution time: 5.004637002944946s

As you see now the time is calculated correctly.

Conclusion

Hope this article could cover some basics and give you the idea behind decorators. Of course you can go as crazy as you want when understanding this concept. E.g. mess around with data returned from a funciton:

# A decorator to convert meters to kilometers
def m_to_km(func):
    return lambda: func() / 1000.

@m_to_km
def get_distance():
    # some complex calculation here...
    return 10000

get_distance()
>>> 10.0

You are limited only by your imagination (and Python compiler) here.