Decorators

Decorators are wrapper functions that enhance the wrapped function. It does not change the behaviour of the original function, but instead wrap it with some new behavior, achieving extensibility and composability.

Basic Example

def timer(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)  # call the function
        total_time = time.time() - start_time
        print(f'execution time: {total_time}')
        return result
    return wrapper

@timer
def slow_func(a, b):
    time.sleep(1)
    return a * b

slow(5,6)
# execution time: 1.001169204711914
# 30

Decorators are however not the only way to transform a function. The same could be achieved like this:

slow = timer(slow)
slow(5,6)

Template

In basic the template for a wrapper is this:

def decorator(func):
     def wrapper(*args, **kwargs):
         # something before
         result = func(*args, **kwargs)
         # something after
         return result

     return wrapper

Decorator Arguments

What if you want to change the output of the decorator function timer? For example, only show the number instead of a pretty text.

Remember that there is also a way to write it, without the decorator:

slow = timer(prettyprint=False)(slow)
slow(5,6)

From that, it follows that the decorator template would then look like this:

def timer(prettyprint=True):
     def wrapped(func):
         def wrapper(*args, **kwargs)
            # code
            return func(*args, **kwargs)
         return wrapper
     return wrapper

The first line receives expects the argument as given in the decorator @timer(prettyprint=False), then wrapped gets the actual function, and in turn calls wrapper() with arguments.

Decorate Class Methods

Classes or methods in a class can have decorators as well. This example adds a reusable function that logs something each time the class is instantiated

import logging
import time
logging.basicConfig(filename='example.log', encoding='utf-8', level=logging.DEBUG)

def logger(cls):
    # do something with class cls here

    def wrapper(*args, **kwargs):
        logging.info(f'{time.time()} {cls.__class__.__name__} ')

        obj = cls(*args, **kwargs)
        # do something with the object here
        return obj

    return wrapper

@logger
class A:
    def __init__(self, x, y):
        self.x = x
        self.y = y

objs = [A(x=1, y=2) for _ in range(10)]

Decorate a Class

Classes can be decorated as well and it gives us some re-usable code that can change the behaviour of a class. It could be used to add or overwrite attributes of classes. The @dataclass decorator is a good example of this.

In this example we will replace the __repr__ method:

def different_repr(cls):
     cls.__repr__ = new_repr  # a new function

     def wrapper(*args, **kwargs):
         return cls(*args, **kwargs)

     return wrapper

functools.wraps

When you use a decorator, you're replacing one function with another. Unfortunately this means when checking __name__, it will print the name of the decorator because that's the name of your new function. functools.wraps takes a function used in a decorator and adds the functionality of copying over the function name, docstring, arguments list, etc.

To complete the timer decorator from the first example, we should add @wraps:

from functools import wraps

def timer(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)  # call the function
        total_time = time.time() - start_time
        print(f'execution time: {total_time}')
        return result
    return wrapper