Python Decorators in Real Code: Timing, Retry and Logging

A decorator is a function that takes a function and returns a new one that wraps it. The @ line is just a shorter way to write that. This lesson builds three decorators you would actually use: timing, retry and call logging, plus the version that takes arguments.

What @ actually does

def loud(func):
    def wrapper(*args, **kwargs):
        print(f"calling {func.__name__}")
        return func(*args, **kwargs)
    return wrapper


@loud
def add(a, b):
    return a + b

# The @loud line above is exactly the same as writing:
# add = loud(add)

print(add(2, 3))

calling add
5

wrapper is a closure. It remembers func, the original add. *args, **kwargs accepts any arguments and passes them on, so the same decorator works on any function. And return func(...) matters: forget it and every decorated function returns None.

Always use functools.wraps

import functools


def loud(func):
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


def loud_fixed(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper


@loud
def add(a, b):
    """Add two numbers."""
    return a + b


@loud_fixed
def subtract(a, b):
    """Subtract b from a."""
    return a - b


print(add.__name__, add.__doc__)
print(subtract.__name__, subtract.__doc__)

wrapper None
subtract Subtract b from a.

Without wraps, every decorated function calls itself wrapper and loses its docstring. Logs, error messages, help() and test reports all show the wrong name. One line fixes it. There is no good reason to leave it out.

Decorator 1: timing

import functools
import time


def timed(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            elapsed = time.perf_counter() - start
            print(f"{func.__name__} took {elapsed:.3f}s")
    return wrapper


@timed
def build_report(rows):
    time.sleep(0.2)  # stands in for real work
    return f"{rows} rows"


print(build_report(500))

build_report took 0.20...s
500 rows

The try/finally means the time is reported even when the function raises an error, which is often exactly when you want to know how long it ran. Use time.perf_counter(), not time.time(), for measuring durations. It is meant for that and isn't affected by the system clock being adjusted.

Decorator 2: retry, with arguments

A retry decorator needs settings: how many attempts, and which errors are worth retrying. A decorator with arguments is one more layer: a function that takes the settings and returns the actual decorator.

import functools
import time


def retry(attempts=3, delay=0.1, on=(ConnectionError,)):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except on as error:
                    if attempt == attempts:
                        raise
                    print(f"{func.__name__} failed ({error}), retry {attempt}/{attempts - 1}")
                    time.sleep(delay * attempt)
        return wrapper
    return decorator


calls = {"n": 0}


@retry(attempts=3, on=(ConnectionError,))
def fetch_rates():
    calls["n"] += 1
    if calls["n"] < 3:
        raise ConnectionError("timeout")
    return {"USD": 56.1}


print(fetch_rates())

fetch_rates failed (timeout), retry 1/2
fetch_rates failed (timeout), retry 2/2
{'USD': 56.1}

The fake fetch_rates fails twice, then succeeds, so you can see the retries. The rate value is made up for the example. Read @retry(attempts=3) as: call retry(attempts=3), which returns decorator, then apply that to fetch_rates. Three levels: settings, function, call.

When the error never goes away, the last attempt re-raises it, so the caller still sees the real error instead of a silent None:

import functools


def retry(attempts=3, on=(ConnectionError,)):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except on:
                    if attempt == attempts:
                        raise
        return wrapper
    return decorator


@retry(attempts=2)
def always_down():
    raise ConnectionError("server unreachable")


try:
    always_down()
except ConnectionError as error:
    print("gave up:", error)

gave up: server unreachable

Decorator 3: call logging

import functools
import logging

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
log = logging.getLogger(__name__)


def logged(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        log.info("%s called with %r %r", func.__name__, args, kwargs)
        try:
            result = func(*args, **kwargs)
        except Exception:
            log.exception("%s failed", func.__name__)
            raise
        log.info("%s returned %r", func.__name__, result)
        return result
    return wrapper


@logged
def apply_discount(amount, percent=0):
    if not 0 <= percent <= 100:
        raise ValueError("bad percent")
    return amount * (100 - percent) / 100


apply_discount(1000, percent=10)
try:
    apply_discount(1000, percent=150)
except ValueError:
    pass

INFO apply_discount called with (1000,) {'percent': 10}
INFO apply_discount returned 900.0
INFO apply_discount called with (1000,) {'percent': 150}
ERROR apply_discount failed
Traceback (most recent call last):
...
ValueError: bad percent

Note the raise after logging. The decorator records the error and lets it continue to the caller. A decorator that logs and swallows errors turns a crash into a wrong result, which is harder to find.

Stacking decorators

You can apply several. They wrap from the bottom up:

@timed
@retry(attempts=3)
def fetch_rates():
    ...

This is timed(retry(attempts=3)(fetch_rates)). timed is outside, so it measures all attempts together, including the waits. Swap the order and it reports each attempt separately. Neither is wrong. Decide which number you actually want.

Common mistakes

• Forgetting return inside the wrapper. The function runs, and every call returns None.
• Writing @retry instead of @retry() for a decorator that takes arguments. Then the function itself is passed as attempts, and you get a confusing error when it is called.
• Doing work at decoration time. Code in the decorator but outside wrapper runs once, at import. Code inside wrapper runs on every call. Put per-call work inside.

What happens if you skip this step

The same timing, retry and logging code gets pasted into every function that needs it, each copy slightly different. One copy retries on every exception, one forgets to re-raise, one logs the password argument. With a decorator, the rule lives in one place, and fixing it fixes it everywhere.

You decorate a function that sends invoices with @retry(attempts=3). The first attempt times out after the server already sent the invoice. What does the customer receive?