Write Typed Decorators with wraps and ParamSpec

You want to print how long each of several functions takes. Instead of adding timing code to each one, you write it once, as a decorator:

import time

def timed(fn):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        print(f"{fn.__name__} took {time.perf_counter() - start:.4f}s")
        return result
    return wrapper

@timed
def add(a: int, b: int) -> int:
    '''Add two numbers.'''
    return a + b

print(add(2, 3))  # prints the time, then 5

A decorator is a function that takes a function and returns a replacement. Writing @timed above def add does the same as writing add = timed(add) after it. The name add now refers to wrapper. The wrapper function collects any arguments in *args and **kwargs, passes them on to the original add, and returns its result.

This works, but it has two problems. Python now reports add.__name__ as "wrapper" and add.__doc__ as None, so logs and help(add) describe the wrapper. And mypy --strict rejects timed, because it has no annotations.

Copy the name and docstring with wraps

Put @functools.wraps(fn) on the wrapper. It copies the original function's __name__, __doc__ and a few other attributes onto the wrapper, so add.__name__ is "add" again and help(add) shows its docstring.

A TypeVar stands for one type

A type variable stands for a type that mypy fills in at each call. You declare one with TypeVar:

from typing import TypeVar

T = TypeVar("T")

def first(items: list[T]) -> T:
    return items[0]

n = first([1, 2])      # mypy reads n as int
s = first(["a", "b"])  # mypy reads s as str

The decorator needs one for its return type. Its wrapper returns whatever the wrapped function returns, and a TypeVar named R says exactly that.

A ParamSpec stands for a parameter list

A ParamSpec is the same idea for a whole list of parameters. After P = ParamSpec("P"), the type Callable[P, R] means "a function that takes parameters P and returns R". Inside the wrapper, *args: P.args and **kwargs: P.kwargs say that the wrapper takes exactly those parameters.

Here is the timing decorator with both, and with wraps:

import time
from collections.abc import Callable
from functools import wraps
from typing import ParamSpec, TypeVar

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

def timed(fn: Callable[P, R]) -> Callable[P, R]:
    @wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        print(f"{fn.__name__} took {time.perf_counter() - start:.4f}s")
        return result
    return wrapper

@timed
def add(a: int, b: int) -> int:
    '''Add two numbers.'''
    return a + b

print(add(2, 3))     # prints the time, then 5
print(add.__name__)  # add

Now mypy knows that the decorated add still takes two int values and returns an int, so it reports add("2", 3) as an error. If you annotate the decorator with Callable[..., Any] instead, mypy accepts any call, and every result has type Any.

The newer spelling

Python 3.12 added a shorter way to declare these. List them in square brackets after the function name, with ** before a ParamSpec, and skip the TypeVar and ParamSpec calls:

from collections.abc import Callable
from functools import wraps

def logged[**P, R](fn: Callable[P, R]) -> Callable[P, R]:
    @wraps(fn)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"calling {fn.__name__}")
        return fn(*args, **kwargs)
    return wrapper

Both spellings mean the same to mypy. The exercises in this course use the older one, which you will also see in most existing code.

Common mistakes

Forgetting to return the result. If the wrapper calls fn but does not return its value, every decorated function returns None. With -> R on the wrapper, mypy reports a missing return.

Forgetting @wraps. Everything still runs, but logs and help() name wrapper instead of your function.

Check your understanding What does @timed above def add do, written without the @? Show answer

It does the same as add = timed(add) right after the def. Python creates the function, passes it to timed, and binds the name add to whatever timed returns, here wrapper. That is why, without wraps, the name of add becomes "wrapper".

Check your understanding You want a decorator that adds a retries=3 keyword to every function it wraps. Can Callable[P, R] describe the result? Show answer

No. The type Callable[P, R] says that the decorated function takes exactly the original's parameters, no more. Put the option on the code that builds the wrapper instead. You can write @retry(times=3), or pass the option to a function that wraps, as make_retrier(fn, attempts=3, ...) does in the next exercise.

Going further

The wraps decorator also sets __wrapped__ on the wrapper to the original function. The call inspect.signature(add) follows that link, so it shows (a: int, b: int) -> int, not (*args, **kwargs). A decorator that adds or removes a leading parameter, such as a database connection, can describe that with typing.Concatenate.