Advanced20 min

Decorators and context managers

Wrap functions with reusable behaviour and manage resources cleanly — the mechanics behind @app.get, @task, and torch.no_grad().

Why this matters in AI / ML / GenAI

FastAPI routes, Airflow tasks, LangChain tools, pytest fixtures, and caching all use decorators. Context managers open and close model sessions, database connections, and inference modes. You will read them daily, and writing a timing or retry decorator saves repetition.

Functions are objects

You can pass a function as an argument, return it from another function, and store it in a dict. A decorator is just a function that takes a function and returns a replacement.

@timed above a definition is shorthand for fn = timed(fn). Nothing magic.

Always apply @functools.wraps(fn) to the inner wrapper so the original name and docstring survive — otherwise every decorated function reports itself as wrapper, which wrecks tracebacks and API docs.

Practical decorators

Three you will actually write:

Timing — log how long an inference call took
Retry — re-run on transient API errors
Caching — @functools.lru_cache memoizes pure functions; great for embedding lookups of repeated strings in a single process

Decorators with arguments need one more layer: a function that returns a decorator (@retry(max_attempts=3)).

Context managers

with calls __enter__ on entry and __exit__ on exit — even when an exception is raised. That is the guarantee that makes with open(...) safe.

Write one quickly with @contextlib.contextmanager: code before yield is setup, code after is teardown, and a try/finally ensures teardown always runs.

Real uses: timing a block, opening a DB session, torch.no_grad() to disable gradient tracking during inference, and temporarily swapping a config value in tests.

Copy-paste examples

Copy into your own editor, or load one into the compiler below and press Run.

A timing decorator

functools.wraps keeps the original function name intact.

import functools
import time

def timed(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        elapsed_ms = (time.perf_counter() - start) * 1000
        print(f"{fn.__name__} took {elapsed_ms:.2f} ms")
        return result
    return wrapper

@timed
def embed(texts):
    return [[len(t) * 0.1] * 3 for t in texts]

print(embed(["python", "genai"]))
print("name preserved:", embed.__name__)

Decorator with arguments (retry)

Three layers: retry -> decorator -> wrapper. Read it from the inside out.

import functools

def retry(max_attempts=3, exceptions=(TimeoutError,)):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            last = None
            for attempt in range(1, max_attempts + 1):
                try:
                    return fn(*args, **kwargs)
                except exceptions as err:
                    last = err
                    print(f"attempt {attempt} failed: {err}")
            raise last
        return wrapper
    return decorator

state = {"calls": 0}

@retry(max_attempts=4)
def call_llm(prompt):
    state["calls"] += 1
    if state["calls"] < 3:
        raise TimeoutError("gateway timeout")
    return f"answer to: {prompt}"

print(call_llm("what is llmops?"))

A context manager for timing a block

Same shape as torch.no_grad() — setup, yield, guaranteed teardown.

import contextlib
import time

@contextlib.contextmanager
def stage(name):
    print(f"-> start {name}")
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed_ms = (time.perf_counter() - start) * 1000
        print(f"<- done {name} in {elapsed_ms:.2f} ms")

with stage("retrieve"):
    docs = [f"doc-{i}" for i in range(1000)]

with stage("generate"):
    answer = " ".join(docs[:3])
print(answer)

Caching repeated work

lru_cache only works when arguments are hashable and the function is pure.

import functools

@functools.lru_cache(maxsize=128)
def fake_embed(text):
    print("computing embedding for:", text)
    return sum(ord(c) for c in text) % 997

print(fake_embed("python"))
print(fake_embed("python"))   # cached, no recompute line
print(fake_embed("genai"))
print(fake_embed.cache_info())

Combine a decorator and a context manager

Try it — in-browser Python

Add a second @timed function and call it inside the with block.

Output

Python runs in your browser. First run downloads the runtime.

Press Run (or Ctrl+Enter) to execute.

CPython in WebAssembly. Stdlib works. NumPy and pandas load on demand. No input(), no GPU, no network installs.

Takeaways

A decorator is a function returning a wrapped function — always use functools.wraps.
Decorators with arguments need an extra layer: retry(...) returns the decorator.
Context managers guarantee teardown; write them with @contextlib.contextmanager.