lesson 162 Decorators & Context Managers · advanced

Decorators with arguments

@retry_on(TimeoutError, attempts=3) is a call that returns a decorator. The outer function receives the settings, builds the decorator, and the wrapper inside can use both the settings and the original function.

You'll be able to

Read @deco(arg) as f = deco(arg)(f) and name the three layers
Write decorators configured by arguments: required keys, value limits, retries
Keep per-function state and injected dependencies (sleep, clock) in the closure

The idea

Many decorators need settings: which exceptions to retry, how many attempts, which keys a record must have. Python handles that with one more layer. In @require_keys("id", "properties") the expression after @ is evaluated first: require_keys("id", "properties") returns a decorator, and that decorator is applied to the function. So @deco(arg) means f = deco(arg)(f).

The code has three nested functions: the factory takes the settings, the decorator takes the function, and the wrapper takes the call's arguments. Each inner function can see the names of the ones around it.

Run each example to see what it prints. Running examples doesn't count toward progress.

example-1.py · Required keys for a record

def require_keys(*keys):
    def decorator(func):
        def wrapper(record, *args, **kwargs):
            missing = [key for key in keys if key not in record]
            if missing:
                raise KeyError(f"{func.__name__} needs {', '.join(missing)}")
            return func(record, *args, **kwargs)
        return wrapper
    return decorator

@require_keys("id", "properties")
def summary(feature):
    return f"{feature['id']}: M{feature['properties']['mag']}"

print(summary({"id": "us6000pgsd", "properties": {"mag": 5.5}}))
try:
    summary({"properties": {"mag": 4.7}})
except KeyError as err:
    print("KeyError:", err)

Settings can be behaviour too. A retry decorator needs to know which exceptions are temporary, how many attempts to make, and how to wait. Passing sleep in as a setting means tests (and this example) can record the waits instead of actually waiting.

example-2.py · Retry chosen exceptions

import time

def retry_on(*exceptions, attempts=3, sleep=time.sleep):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions:
                    if attempt == attempts:
                        raise
                    sleep(2 ** (attempt - 1))
        return wrapper
    return decorator

waits = []
outcomes = iter([ConnectionError("reset"), TimeoutError("slow"), {"count": 7}])

@retry_on(ConnectionError, TimeoutError, attempts=4, sleep=waits.append)
def fetch_count():
    outcome = next(outcomes)
    if isinstance(outcome, Exception):
        raise outcome
    return outcome

print(fetch_count(), waits)

When to use it: whenever a decorator's behaviour should vary per function: limits, names, exception types, intervals. Mistakes to avoid: writing @retry_on without the brackets for a factory (the function is passed as the first *setting*, and the result isn't a working decorator); returning the wrong layer (each layer must return the one inside it); and putting per-call state in the factory, where every call of the decorated function shares it.

Review

@deco(arg) means f = deco(arg)(f): a factory returns the decorator.
Three layers: settings, function, call. Each returns the one inside it.
Settings can be behaviour, such as sleep or clock, which keeps the decorator testable.
Next: functools.wraps, so a wrapper keeps the original's name and documentation.