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.