Decorators

TL;DR

A decorator is a function that takes a function and returns a function. The @decorator syntax above a def is sugar for fn = decorator(fn). Everything else — the wrapped behaviour, the arguments, the registration patterns — is built on that one rule.

In real ML codebases, decorators are how you add cross-cutting concerns without polluting every function: caching with @functools.cache, disabling gradients with @torch.no_grad, registering FastAPI routes with @app.post(...), retrying flaky LLM calls, timing slow code, validating inputs. They show up everywhere; understanding the shape makes a lot of library magic stop feeling like magic.

@decorator
def fn(x): ...

# is exactly:
def fn(x): ...
fn = decorator(fn)

That’s the whole syntax. The rest is just choosing what decorator does.

The picture in your head

A decorator slips a wrapper around your function. The caller still calls fn(x, y), but what actually runs is wrapper(x, y), which can do work before, after, instead of, or around the original call.

def trace(fn):
    def wrapper(*args, **kwargs):
        print(f"calling {fn.__name__}({args}, {kwargs})")
        result = fn(*args, **kwargs)
        print(f"  -> {result}")
        return result
    return wrapper

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

>>> add(2, 3)
calling add((2, 3), {})
  -> 5
5

trace took the function add and returned wrapper, which now replaces add in the namespace. Calling add(2, 3) calls wrapper(2, 3), which calls the real add in the middle.

The minimal correct decorator

The above trace has two bugs that bite real codebases:

1. wrapper.__name__ is now "wrapper", not "add". Stack traces and help(add) lie.
2. Decorating a method on a class breaks introspection downstream.

Fix both with functools.wraps:

import functools

def trace(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        print(f"calling {fn.__name__}")
        return fn(*args, **kwargs)
    return wrapper

@functools.wraps(fn) copies __name__, __doc__, __module__, __qualname__, and __wrapped__ onto the wrapper. Always use it. The omission is the #1 sign of a homemade decorator written in a hurry.

Decorators that take arguments

Sometimes you want @retry(max_attempts=3), not just @retry. The trick: the outer function takes the decorator’s arguments, returns the actual decorator.

import functools, time, random

def retry(max_attempts=3, backoff=0.5):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            for attempt in range(1, max_attempts + 1):
                try:
                    return fn(*args, **kwargs)
                except Exception as e:
                    if attempt == max_attempts:
                        raise
                    time.sleep(backoff * attempt)
            # unreachable
        return wrapper
    return decorator

@retry(max_attempts=5, backoff=1.0)
def call_llm(prompt: str) -> str:
    return openai_client.complete(prompt)

The shape is always: decorator_factory(args) -> decorator -> wrapper. Three nested functions.

For decorators that might take arguments, the canonical idiom is to inspect the first argument:

def cached(maxsize=128):
    if callable(maxsize):                  # used as @cached without ()
        return functools.lru_cache()(maxsize)
    return functools.lru_cache(maxsize=maxsize)

This lets @cached and @cached(maxsize=64) both work. Used by many real libraries (pytest.fixture, for instance).

The decorators you’ll actually use

A working ML engineer reaches for these constantly.

functools.cache — memoise

import functools

@functools.cache
def expensive(model_name: str) -> Model:
    return load_from_disk(model_name)

expensive("gpt-4")    # first call: loads from disk
expensive("gpt-4")    # second call: instant, returns cached object

cache is unbounded (3.9+). Use lru_cache(maxsize=N) if you need a ceiling.

The arguments must be hashable — cache keys on the args. Lists and dicts can’t be cached on; tuples and strings can.

functools.lru_cache — bounded memoise

@functools.lru_cache(maxsize=10_000)
def embed(text: str) -> tuple[float, ...]:
    return tuple(model.encode(text))   # tuple, not ndarray, so it's hashable

LRU evicts the least recently used when the cache is full.

@property — attribute-style methods

class Model:
    def __init__(self, weights):
        self._weights = weights

    @property
    def num_params(self) -> int:
        return sum(w.numel() for w in self._weights)

>>> m = Model([...])
>>> m.num_params         # called like an attribute, not m.num_params()
1234567

@staticmethod and @classmethod

class Tokenizer:
    def __init__(self, vocab): ...

    @classmethod
    def from_file(cls, path: str) -> "Tokenizer":
        return cls(load_vocab(path))

    @staticmethod
    def is_special_token(tok: str) -> bool:
        return tok.startswith("<") and tok.endswith(">")

>>> Tokenizer.from_file("vocab.txt")    # alternative constructor
>>> Tokenizer.is_special_token("<bos>") # utility, no instance needed

@dataclass — see PY 108

The whole dataclasses module is decorator-driven.

Library-specific essentials

• PyTorch: @torch.no_grad(), @torch.inference_mode() — disable autograd inside an inference function.
• FastAPI: @app.get("/path"), @app.post(...) — route registration.
• pytest: @pytest.fixture, @pytest.mark.parametrize(...), @pytest.mark.skip — test machinery.
• Click: @click.command(), @click.option(...) — CLI building.
• Numba: @numba.jit, @numba.njit — JIT-compile a Python function to LLVM-compiled native code.

Stacking decorators

Multiple decorators stack bottom-up — the one closest to the function runs first.

@cache
@retry(max_attempts=3)
def call_llm(prompt):
    ...

Reading order: call_llm = cache(retry(max_attempts=3)(call_llm)). So retry wraps call_llm first; cache wraps the retrying version. Order matters: this caches successful results (good), but if you put them the other way around — @retry outside @cache — you’d retry on cache misses, which is also fine but means the cache is bypassed during retries. Think about it.

Class decorators

A decorator can be applied to a class too. It receives the class and returns a class.

def register(cls):
    REGISTRY[cls.__name__] = cls
    return cls

@register
class CrossEntropyLoss: ...

@register
class FocalLoss: ...

# Now REGISTRY = {"CrossEntropyLoss": ..., "FocalLoss": ...}

This is the registry pattern, used by Hugging Face, MMDetection, and most other plugin-style ML libraries.

@dataclass is a class decorator that synthesises __init__, __repr__, __eq__. It’s the most-used class decorator in the ecosystem.

A worked example — a timer that logs

A useful decorator I add to most ML codebases on day one.

import functools
import logging
import time
from typing import Callable, ParamSpec, TypeVar

logger = logging.getLogger(__name__)

P = ParamSpec("P")
R = TypeVar("R")

def timed(fn: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        t0 = time.perf_counter()
        try:
            return fn(*args, **kwargs)
        finally:
            elapsed = time.perf_counter() - t0
            logger.info("%s took %.3fs", fn.__name__, elapsed)
    return wrapper

@timed
def train_epoch(loader, model, optimizer):
    for batch in loader:
        ...

Note ParamSpec and TypeVar — these preserve the wrapped function’s type signature so static type-checkers see train_epoch as having its real signature, not (*args, **kwargs). See PY 107.

Common gotchas

• Forgetting functools.wraps. Stack traces lose the function’s real name. inspect.signature lies. Use it always.
• Decorator runs at import time. @register adds the class to the registry the moment the module is imported. If your registry is empty, suspect the module hasn’t been imported.
• Decorators with arguments add a layer. @retry (no parens) treats the function as max_attempts. @retry() is the parameterised version. Mixing them up is a common bug.
• @cache on a method caches self. The instance is part of the hash key. If self is unhashable (most class instances are by default), it crashes. Either define __hash__, use cached_property, or factor the cached logic into a module-level function.
• @cache keeps strong references forever. It’s a memory leak waiting to happen for short-lived objects passed in as args.
• Decorators on async functions need to be async-aware. @retry written for sync functions will return a coroutine without awaiting it. Need a separate async version, or detect with inspect.iscoroutinefunction.

When NOT to use a decorator

• One-off setup that’s only used once. Just call the function. A decorator earns its keep when applied many times.
• Logic that needs to inspect or modify the function’s body. You can’t. Decorators wrap; they don’t transform source. (For source transformation, you want AST manipulation or a code-gen tool.)
• Behaviour that should be configurable per-call. A decorator is baked at function-definition time. If you need “sometimes retry, sometimes don’t,” pass a parameter instead.

Where decorators show up in real ML codebases

• @torch.no_grad() wraps every inference function.
• @functools.cache memoises tokenizer lookups, model loads, embedding computations.
• @app.post("/predict") is how FastAPI knows your function is a route.
• @pytest.fixture is the entire fixture system.
• @dataclass for every config and result type.
• @register("loss") patterns in Hugging Face / MMDetection / Lightning.
• @hydra.main for CLI / config-driven entry points in Hydra-based projects.
• @retry (from tenacity) for every external API call.
• @profile from line_profiler for finding slow lines.

The defensive habit: when you find yourself adding the same five lines of setup/cleanup at the top and bottom of multiple functions, that’s a decorator waiting to be written.

Resources

• Python docs — functools — docs.python.org — cache, lru_cache, wraps, partial, reduce.
• PEP 318 — Decorators for Functions and Methods — peps.python.org — the original spec.
• PEP 3129 — Class Decorators — peps.python.org.
• Real Python — Primer on Decorators — realpython.com — long, thorough, mostly correct.
• tenacity — tenacity.readthedocs.io — the production-grade retry decorator. Pip-install it instead of writing your own.