14. Decorators

14.1 A decorator is a function that wraps a function

import functools
import time

def timed(fn):
    @functools.wraps(fn)                  # copy __name__, __doc__, etc. to the wrapper
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        try:
            return fn(*args, **kwargs)
        finally:
            print(f"{fn.__name__} took {(time.perf_counter() - start) * 1000:.2f} ms")
    return wrapper

@timed                                    # same as: slow_sum = timed(slow_sum)
def slow_sum(n: int) -> int:
    return sum(range(n))

print(slow_sum(1_000_000))
print(slow_sum.__name__)                  # slow_sum (thanks to functools.wraps)

@decorator above a def is only syntax for f = decorator(f). Because functions are objects (Ch 5.2) and closures capture variables (Ch 5.4), a decorator can add behaviour before and after any call without changing the function. Frameworks use this everywhere: @app.get("/users") in FastAPI, @pytest.fixture, @torch.no_grad(), @functools.cache.

14.2 Decorators with arguments: a retry decorator

A decorator that takes arguments is a function that returns a decorator. This one retries a flaky operation with exponential backoff and jitter, the policy from the HTTPS guide Ch 21:

import functools
import random
import time

def retry(*, attempts: int = 3, base: float = 0.1, on: tuple[type[Exception], ...] = (OSError,)):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return fn(*args, **kwargs)
                except on as e:
                    if attempt == attempts:
                        raise
                    delay = base * 2 ** (attempt - 1) * random.uniform(0.5, 1.5)
                    print(f"{fn.__name__}: {e!r}; retry {attempt}/{attempts - 1} in {delay:.2f}s")
                    time.sleep(delay)
        return wrapper
    return decorator

calls = 0

@retry(attempts=4, base=0.05, on=(ConnectionError,))
def flaky() -> str:
    global calls
    calls += 1
    if calls < 3:
        raise ConnectionError("reset by peer")
    return "ok"

print(flaky(), "after", calls, "calls")

Only retry idempotent operations. Retrying a POST /payments that timed out after the server processed it charges the customer twice. Ch 43 adds idempotency keys.

14.3 Caching with functools

import functools

@functools.cache                    # unbounded memoization
def fib(n: int) -> int:
    return n if n < 2 else fib(n - 1) + fib(n - 2)

print(fib(200))                     # instant; exponential without the cache

@functools.lru_cache(maxsize=1024)  # bounded: evicts least-recently used
def resolve(host: str) -> str:
    return f"10.0.0.{len(host)}"    # pretend this is a slow DNS lookup

resolve("api"); resolve("api")
print(resolve.cache_info())         # CacheInfo(hits=1, misses=1, maxsize=1024, currsize=1)

Cache keys are the arguments, so they must be hashable. A cache on a method (@cache def m(self, ...)) keeps every self alive forever, a memory leak; use functools.cached_property for per-instance values instead. A cache also needs an invalidation story: the DNS cache above would serve a stale address after a failover, which is exactly the DNS TTL problem.