typing decorators

Python decorators are not the easiest to write type annotations for. So normally, I’d just skip the annotations altogether. As an example, here’s a decorator that prints out debug information:

def debug(func):
    def wrapper(*args, **kwargs):
        print(f"Called {func.__name__!r} args={args} kwargs={kwargs}")
        result = func(*args, **kwargs)
        print(f"Result: {result} {type(result)}")
        return result

    return wrapper

which you could then apply to a function:

@debug
def add(x: int, y: int) -> int:
    return x + y

add(2, 4)
# Called 'add' args=(2, 4) kwargs={}
# Result: 6 <class 'int'>

Everything works as it should, but the lack of types makes it difficult for consumers of the code. So, let’s add some types that encode the following information:

1. the returned function should share the same signature as the one passed.
2. the wrapped function may vary, so its signature must be generic.

Python’s new generic typing syntax simplifies this for us. At the time of writing, I didn’t have a solution for decorators, but now I do:

from collections.abc import Callable

def debug[**P, R](func: Callable[P, R]) -> Callable[P, R]:
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"Called {func.__name__!r} args={args} kwargs={kwargs}")
        result = func(*args, **kwargs)
        print(f"Result: {result} {type(result)}")
        return result

    return wrapper

Let’s unpack this change:

• [**P, R] is shorthand for generic ParamSpec and TypeVar variables.
• Callable[P, R], is what the decorator (wrapper) receives and returns.
• Lastly, the args and kwargs components of the parameter specification P is used to type the wrapper function.

Currently, the decorator accepts any function no matter what the arguments or return value are. This fine in our current use case, but say we want to perform an action on the result which is only possible for certain types.

For example, if you wanted to add another debug utility that prints out the length of the return value. You could try and use the same approach as before:

from collections.abc import Callable

def debug_length[**P, R](func: Callable[P, R]) -> Callable[P, R]:
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        result = func(*args, **kwargs)
        print(f"Length of result: {len(result)}")
        return result

    return wrapper

However, you’d get the following error:

Argument 1 to "len" has incompatible type "R"; expected "Sized"  [arg-type]

This error makes sense. As you can’t call len with just any type. Python luckily provides an abstract type called Sized to help us (which mypy rightfully points out). With this type, we can now bind our generic return type to this interface:

from collections.abc import Callable, Sized

def debug_length[**P, R: Sized](func: Callable[P, R]) -> Callable[P, R]:
    ...

With this change, we are no longer getting errors when defining the decorator. But perhaps more importantly, we now get an error when applying it to a function with an unsupported return type:

@debug_length
def add(x: int, y: int) -> int:
    return x + y

Which results in the following error:

error: Value of type variable "R" of "debug_length" cannot be "int"  [type-var]

Whereas, the functions below are still considered valid:

@debug_length
def get_letters() -> str:
    return string.ascii_lowercase

@debug_length
def multiply_items(nums: list[int], *, multiplier: int) -> list[int]:
    return [n * multiplier for n in nums]

What’s great about this is that decorator is still quite generic. Your function can still receive any argument. The key aspect is that the return types, although different, both adhere to the Sized protocol.

Your colleague is grateful for all these typing tips, but is still uncertain:

I mean this is fine for simple decorators, but what if I want to configure it with parameters? For example, I’d like the ability to hide the result output for the debug decorator. Sometimes the result is simply too large and shouldn’t be printed. Do you think you could help me out here?

Luckily, adding parameters is not too difficult (albeit verbose):

def debug[**P, R](
    show_result: bool
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    # notice the additional nesting: debug > decorator > wrapper
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            print(f"Called {func.__name__!r} args={args} kwargs={kwargs}")
            result = func(*args, **kwargs)
            if show_result:
                print(f"Result: {result} {type(result)}")
            return result

        return wrapper

    return decorator

This allows us now to pass explicit flag to hide result output:

@debug(show_result=False)
def add(x: int, y: int) -> int:
    return x + y

add(2, 4)
# Called 'add' args=(2, 4) kwargs={}

This works well, but my goodness, there are a lot of Callable types! In my opinion, this is difficult to read, but there may be a better way. Anthony Sottile introduced a cool alternative in his typing decorators sucks! #573 video. Where he used a contextmanager as a decorator to monitor how long a function takes to run:

import contextlib
import time

@contextlib.contextmanager
def timeit(name: str) -> Generator[None]:
    start = time.perf_counter()
    try:
        yield
    finally:
        end = time.perf_counter()
        print(f"LOG {name} took: {end - start}")

Caveat: you can only do this if you don’t need access to the wrapped function’s result. Whatever happens in the yield block is completely detached from the context manager’s body. If you do need the result, you’re better off writing a traditional decorator like the one above.

Surprisingly, the contextmanager supports the decorator syntax:

@timeit(name="decorator")
def add(x: int, y: int) -> int:
    return x + y

add(2, 4)
# LOG decorator took: 3.167300019413233e-05

But you can (of course) also call it like a traditional context manager:

with timeit("contextmanager"):
    add(2, 4)
# LOG contextmanager took: 3.167300019413233e-05

For comparison, here is the equivalent decorator with the classic approach:

import time

def timeit[**P, R](
    name: str,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            start = time.perf_counter()
            result = func(*args, **kwargs)
            end = time.perf_counter()
            print(f"LOG {name} took: {end - start}")
            return result

        return wrapper

    return decorator

Personally, I find the context manager approach easier on the eyes.

In this article, I’ve shown ways to write type annotations for your decorators. I think for simple decorators the new generic syntax should meet your needs. However, with parameterized decorators, the type signature may become unwieldy. In those cases, you may want to consider a contextmanager (if you don’t need access to the inner result).

Anyway, that’s all I got. If you any additional tips or discovered an error in the post, please send me feedback. Thanks for reading and happy typing!