Decorators are a powerful feature in Python that can help simplify code and improve maintainability. In this article, we will explore how decorators can be used to avoid duplicated code and make it easier to modify functionality across multiple functions.

The Problem: Duplicated Code

Let’s consider an example where we have two functions, add and multiply, which performs some calculations and also logs the execution time.

import time

def add(num1: int, num2: int):
    """Takes two integers and returns their sum."""
    start = time.time()
    print(f"Add {num1} and {num2}")
    res = num1 + num2
    end = time.time()
    print(f'Elapsed time: {(end - start) * 1000:.3f}ms')
    return res

def multiply(num1: int, num2: int):
    """Takes two integers and returns their product."""
    start = time.time()
    print(f"Multiply {num1} and {num2}")
    res = num1 * num2
    end = time.time()
    print(f'Elapsed time: {(end - start) * 1000:.3f}ms')
    return res

add(1, 2)
multiply(1, 2)

Output:

Add 1 and 2
Elapsed time: 0.142ms
Multiply 1 and 2
Elapsed time: 0.006ms
2

As we can see, the timing logic is duplicated in both functions. If we need to modify the timing logic, we would have to update it in both places, which can be time-consuming and error-prone.

The Solution: Decorators

Decorators can help us avoid duplicated code and make it easier to modify functionality across multiple functions.

Let’s define a decorator time_func that can be used to track the execution time of any function.

import time

def time_func(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print(f'Elapsed time: {(end - start) * 1000:.3f}ms')
        return res
    return wrapper

Here’s a step-by-step explanation of the decorator:

1. time_func takes a function func as an argument.
2. The inner wrapper function handles any number of positional and keyword arguments.
3. It records the start time, calls the original function with the original arguments, and records the end time.
4. The elapsed time is calculated and printed.
5. The result of the original function is returned.

We can now apply the time_func decorator to our add and multiply functions.

@time_func
def add(num1: int, num2: int):
    """Takes two integers and returns their sum."""
    print(f"Add {num1} and {num2}")
    return num1 + num2

@time_func
def multiply(num1: int, num2: int):
    """Takes two integers and returns their product."""
    print(f"Multiply {num1} and {num2}")
    return num1 * num2

add(1, 2)
multiply(1, 2)

Output:

Add 1 and 2
Elapsed time: 0.453ms
Multiply 1 and 2
Elapsed time: 0.007ms
2

As we can see, the timing logic is now defined in one place, and we can easily modify it if needed.

Preserving Function Metadata

When we use the time_func decorator, it changes the function name and docstring. To preserve the original function metadata, we can use the wraps decorator from the functools module.

import time
from functools import wraps

def time_func_with_wraps(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print(f'Elapsed time: {(end - start) * 1000:.3f}ms')
        return res
    return wrapper

We can now apply the time_func_with_wraps decorator to our add and multiply functions.

@time_func_with_wraps
def add(num1: int, num2: int):
    """Takes two integers and returns their sum."""
    print(f"Add {num1} and {num2}")
    return num1 + num2

@time_func_with_wraps
def multiply(num1: int, num2: int):
    """Takes two integers and returns their product."""
    print(f"Multiply {num1} and {num2}")
    return num1 * num2

add(1, 2)
multiply(1, 2)

Output:

Add 1 and 2
Elapsed time: 0.332ms
Multiply 1 and 2
Elapsed time: 0.008ms
2

As we can see, the function name and docstring are now preserved.

print(f"Function name: {add.__name__}")
print(f"Docstring: {add.__doc__}")

Output:

Function name: add
Docstring: Takes two integers and returns their sum.