Decorators & Context Managers

Functions in Python are objects: you can pass them around and return them. A decorator is a function that takes a function and returns an improved version — adding timing, logging, caching, retries or permission checks without touching the original code. Always use functools.wraps so the wrapped function keeps its name and docstring.

Decorators can take arguments too (like @retry(times=3)); that just adds one more layer of function.

A context manager sets something up and guarantees cleanup — files, locks, database transactions, timers. Write one quickly with @contextmanager and a single yield, placing cleanup in finally.

decorators.pyPython

import time
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from functools import lru_cache, wraps
from typing import ParamSpec, TypeVar

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


def timed(func: Callable[P, R]) -> Callable[P, R]:
    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        start = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            print(f"{func.__name__} took {(time.perf_counter() - start) * 1000:.1f} ms")
    return wrapper


def retry(times: int) -> Callable[[Callable[P, R]], Callable[P, R]]:
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except ConnectionError:
                    if attempt == times:
                        raise
                    print(f"attempt {attempt} failed, retrying")
            raise AssertionError("unreachable")
        return wrapper
    return decorator


@contextmanager
def section(title: str) -> Iterator[None]:
    print(f"--- {title} ---")
    try:
        yield
    finally:
        print(f"--- end {title} ---")


@lru_cache(maxsize=None)
def fib(n: int) -> int:
    return n if n < 2 else fib(n - 1) + fib(n - 2)


calls = {"n": 0}


@retry(times=3)
def flaky() -> str:
    calls["n"] += 1
    if calls["n"] < 3:
        raise ConnectionError("network down")
    return "ok"


@timed
def slow_sum(n: int) -> int:
    return sum(range(n))


with section("demo"):
    print(slow_sum(1_000_000))
    print(fib(80))          # instant thanks to lru_cache
    print(flaky())

Runs in your browser · Python

Key points

A decorator = a function that wraps another function; use functools.wraps.
functools.lru_cache is a ready-made memoisation decorator.
@contextmanager + try/finally guarantees cleanup.

Exercise

Write a @require_role("admin") decorator that checks a user dict's role before calling the function and raises PermissionError otherwise. Then write a timer() context manager that prints how long its block took.

Show solution

Try the exercise yourself first — then compare your approach with this one.

require_role takes an argument, so it is a function that returns the real decorator. The wrapper checks the user argument before calling the original function. timer is a generator-based context manager: the code before yield runs on entry, and finally runs on exit even if the block raises.

roles.pyPython

import time
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from functools import wraps
from typing import Any


def require_role(role: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(user: dict[str, str], *args: Any, **kwargs: Any) -> Any:
            if user.get("role") != role:
                raise PermissionError(f"{func.__name__} requires role {role!r}")
            return func(user, *args, **kwargs)
        return wrapper
    return decorator


@contextmanager
def timer(label: str) -> Iterator[None]:
    start = time.perf_counter()
    try:
        yield
    finally:
        print(f"{label} took {(time.perf_counter() - start) * 1000:.1f} ms")


@require_role("admin")
def delete_results(user: dict[str, str], term: int) -> str:
    return f"{user['name']} deleted term {term} results"


with timer("admin action"):
    print(delete_results({"name": "Head Teacher", "role": "admin"}, 2))

try:
    delete_results({"name": "Juma", "role": "student"}, 2)
except PermissionError as err:
    print("Denied:", err)
print(delete_results.__name__)   # delete_results (kept by functools.wraps)

Runs in your browser · Python

Key points

Exercise

Check your understanding