Decorating functions

Decorators are functions that allow us to modify the behavior of other functions. If you have seen a function preceded by @, that is a decorator.

Let’s see an example. We are going to create a decorator that shows the time it took to execute a function. The decorator is defined with def, as if it were a normal function.

The decorator accepts as input a func function and returns it with slightly modified behavior.

import time

# This is a decorator
def measure_time(func):
    def wrapper(*args, **kwargs):
        t0 = time.time()
        res = func(*args, **kwargs)
        print(f"{func.__name__}: {time.time() - t0:.5f} seconds")
        return res
    return wrapper

Now we can use our decorator with @measure_time on any function. For example, the following function determines whether a number is prime or not.

The only change is that it now prints the time it took to run. The rest is the same.

# This is using a decorator
@measure_time
def is_prime(n):
    return n > 1 and all(n % i != 0 for i in range(2, int(n**0.5) + 1))

print(is_prime(99999999999973))
# is_prime: 0.56682 seconds
# True

That is, we have modified the behavior of the function without actually changing it. We have decorated it.

You can also modify the function to return other arguments. In this case, we return the time it took to execute as well as whether it is prime or not.

def measure_time(func):
    def wrapper(*args, **kwargs):
        t0 = time.time()
        # We return the result
        # and the time it took
        return func(*args, **kwargs), time.time() - t0
    return wrapper

Although both examples are valid, there is an important nuance. If we look at the function metadata, it refers to the decorator, not the decorator of our function. For example.

# the decorator has changed the metadata
print(is_prime.__name__)
# wrapper

If we want to preserve metadata, such as name and documentation, a good practice is to use wraps as follows. It is common to see decorators defined in this way.

from functools import wraps
import time

def measure_time(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        t0 = time.time()
        return func(*args, **kwargs), time.time() - t0
    return wrapper

If you now access the name, you can see how it is maintained.

# With wraps now the metadata does not change
print(is_prime.__name__)
# is_prime

On the other hand, it is also possible to pass parameters to our decorators. Just as a function can have input arguments and default values.

Let’s see an example with a decorator that tries to execute a function multiple times in case it fails. We have two parameters:

• 🔁 retries: Defines the number of times the function is attempted to be executed. After this number of attempts, it is considered to have failed.
• 💤 backoff_seconds: Waiting time between consecutive attempts. It is important to define this because if a call fails, if we try without waiting, we may encounter the same error. Waiting a bit between attempts usually increases the chances of success.

This decorator can be useful when we have functions that can sometimes fail. It will allow us to try again when it fails up to a maximum number of times.

def retry(retries, backoff_seconds=1):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, retries + 1):
                try:
                    print(f"Attempt {attempt}")
                    return func(*args, **kwargs)
                except Exception as e:
                    print(f"Error: {e}")
                    if attempt == retries:
                        raise e
                    time.sleep(backoff_seconds)
        return wrapper
    return decorator

Now we can use our decorator in the following way. With random, we try to reproduce a function that can fail with a probability of 30%. Imagine this is a request to an external server or database.

You can see how if failure occurs, it tries again up to a maximum number of 3 times, waiting 2 seconds between attempts.

@retry(retries=3, backoff_seconds=2)
def function_fail():
    import random
    if random.random() < 0.7:
        raise ValueError("Random error")
    print("Completed")