Fully-typed Python decorator with optional arguments#

This is my very first blog post! In this short post, I will show the blueprint for a fully-typed Python decorator that optionally accepts parameters.

The code shown in this short article is taken from my small open source project design-by-contract which provides a typed decorator. Decorators are a very useful concept and you surely will find a lot of introductions to them on the net. In short, they allow executing code every time (before and after) a decorated function is called. That way you can modify function arguments or the returned values, measure time of execution, add logging, perform execution-time type checking and much more. Note that decorators can be written for classes too, offering an alternative method of meta programming (for instance as it is done in the attrs package)

In its simplest form, a decorator is defined similarly to the code below:

def my_first_decorator(func):
    def wrapped(*args, **kwargs):
        # do something before
        result = func(*args, **kwargs)
        # do something after
        return result

    return wrapped

@my_first_decorator
def func(a):
    return a

This works because when the nested function wrapped is defined, its surrounding variables are accessible within in the function and kept in memory as long as the function is used somewhere (this is called a closure in functional programming languages).

Simple enough. However, this has a few disadvantages. The most problematic is that the decorated function will loose its signature (you can see this with inspect.signature), its docstring and even its name! These are problems for source code documentation tools such as sphinx but can easily be resolved with the functools.wraps decorator from the standard library:

from functools import wraps
from typing import Any, Callable, TypeVar, ParamSpec

P = ParamSpec("P") # requires python >= 3.10
R = TypeVar("R")

def my_second_decorator(func: Callable[P, R]) -> Callable[P, R]:
    @wraps(func)
    def wrapped(*args: Any, **kwargs: Any) -> R:
        # do something before
        result = func(*args, **kwargs)
        # do something after
        return result

    return wrapped

@my_second_decorator
def func2(a: int) -> int:
    """Does nothing"""
    return a

print(func2.__name__)
# 'func2'
print(func2.__doc__)
# 'Does nothing'

In this example I already added type annotations. I will elaborate on this in future blog posts but annotations and type hints are the most important additions ever made to Python in my opinion. Better readability, code completion in IDEs, and maintainability of larger code bases are just a few examples of the benefits. The code above should already cover most use cases but it is not possible to parameterize the decorator. Think of writing a decorator that logs the execution time of a function but only if it exceeds a certain amount of seconds. This amount should be configurable for each decorated function separately. If none is specified a default should be used and the decorator should be used without parentheses such that it is easier to use:

@time(threshold=2)
def func1(a):
    ...

# No paranthesis when using default threshold
@time
def func2(b):
    ...

If you can live with parentheses in the second case or don’t offer default values for the parameters at all, then this recipe should be sufficient:

from functools import wraps
from typing import Any, Callable, TypeVar, ParamSpec

P = ParamSpec("P") # requires python >= 3.10
R = TypeVar("R")

def my_third_decorator(threshold: int = 1) -> Callable[[Callable[P, R]], Callable[P, R]]:
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> R:
            # do something before you can use `threshold`
            result = func(*args, **kwargs)
            # do something after
            return result
        return wrapper
    return decorator


@my_third_decorator(threshold=2)
def func3a(a: int) -> None:
    ...

# works
@my_third_decorator()
def func3b(a: int) -> None:
    ...

# Does not work!
@my_third_decorator
def func3c(a: int) -> None:
    ...

To cover the third case, there are some packages for that, namely wrapt and decorator which can actually do much more than just adding optional parameters. While being of very high quality, they introduce a fair deal of additional complexity. With wrapt-decorated functions, I further ran into issues with serialization when running functions on a remote cluster. Both are to my knowledge not fully typed such that static type checkers / linters such as mypy fail in strict mode.

I had to resolve these issues when I was working on my own package and decided to write my own solution. It turned into a pattern that can be easily reused but is difficult to be converted into a library unfortunately.

It uses the overload decorator of the standard library. That way, the same decorator can be specified for use with our without parameters. Apart from that, it is a combination of the two snippets above. One drawback of this approach is that all parameters are required to be given as keyword parameters (which, however, increases readability after all)

from typing import Callable, TypeVar, ParamSpec
from functools import partial, wraps

P = ParamSpec("P") # requires python >= 3.10
R = TypeVar("R")

@overload
def typed_decorator(func: Callable[P, R]) -> Callable[P, R]:
    ...


@overload
def typed_decorator(*, first: str = "x", second: bool = True) -> Callable[[Callable[P, R]], Callable[P, R]]:
    ...


def typed_decorator(
    func: Optional[Callable[P, R]] = None, *, first: str = "x", second: bool = True
) -> Union[Callable[[Callable[P, R]], Callable[P, R]], Callable[P, R]]:
    """
    Describe what the decorator is supposed to do!

    Parameters
    ----------
    first : str, optional
        First argument, by default "x".
        This is a keyword-only argument!
    second : bool, optional
        Second argument, by default True.
        This is a keyword-only argument!
    """

    def wrapper(func: Callable[P, R], *args: Any, **kw: Any) -> R:
        """The actual logic"""
        # Do something with first and second and produce a `result` of type `R`
        return result

    # Without arguments `func` is passed directly to the decorator
    if func is not None:
        if not callable(func):
            raise TypeError("Not a callable. Did you use a non-keyword argument?")
        return wraps(func)(partial(wrapper, func))

    # With arguments, we need to return a function that accepts the function
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        return wraps(func)(partial(wrapper, func))
    return decorator

Later, we can use the decorator with our without arguments respectively

@typed_decorator
def spam(a: int) -> int:
    return a

@typed_decorator(first = "y)
def eggs(a: int) -> int:
    return a

There surely is some overhead to this pattern but the benefit outweighs the costs, don’t you agree?

Updated on 15 September 2022

Fix typos in last example.