How can I type hint decorator that is meant to be used for both sync & async functions?
I've tried something like below, but mypy raises errors:
x/decorator.py:130: error: Incompatible types in "await" (actual type "Union[Awaitable[Any], R]", expected type "Awaitable[Any]") [misc]
x/decorator.py:136: error: Incompatible return value type (got "Union[Awaitable[Any], R]", expected "R") [return-value]
def log_execution_time(foo: Callable[P, AR | R]) -> Callable[P, AR | R]:
module: Any = inspect.getmodule(foo)
module_spec: Any = module.__spec__ if module else None
module_name: str = module_spec.name if module_spec else foo.__module__ # noqa
@contextmanager
def log_timing():
start = time()
try:
yield
finally:
exec_time_ms = (time() - start) * 1000
STATS_CLIENT.timing(
metric_key.FUNCTION_TIMING.format(module_name, foo.__name__),
exec_time_ms,
)
async def async_inner(*args: P.args, **kwargs: P.kwargs) -> R:
with log_timing():
result = await foo(*args, **kwargs) <- error
return result
def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
with log_timing():
result = foo(*args, **kwargs)
return result <- error
if inspect.iscoroutinefunction(foo):
return wraps(foo)(async_inner)
return wraps(foo)(sync_inner)
I know there's a trick like this:
if inspect.iscoroutinefunction(foo):
async_inner: foo # type: ignore[no-redef, valid-type]
return wraps(foo)(async_inner)
sync_inner: foo # type: ignore[no-redef, valid-type]
return wraps(foo)(sync_inner)
But I was hoping that there's a way to properly type hint this.
I'm on python 3.10.10.
PS. I forgot to say that it's important that PyCharm picks it up & suggests proper types.
Say you want to write a function decorator that performs some actions before and/or after the actual function call. Let's call that surrounding context my_context. If you want the decorator to be applicable to both asynchronous and regular functions, you'll need to accommodate both types in it.
How can we properly annotate it to ensure type safety and consistency, while also retaining all possible type information of the wrapped function?
Here is the cleanest solution I could come up with:
from collections.abc import Awaitable, Callable
from functools import wraps
from inspect import iscoroutinefunction
from typing import ParamSpec, TypeVar, cast, overload
P = ParamSpec("P")
R = TypeVar("R")
@overload
def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]: ...
@overload
def decorator(func: Callable[P, R]) -> Callable[P, R]: ...
def decorator(func: Callable[P, R]) -> Callable[P, R] | Callable[P, Awaitable[R]]:
if iscoroutinefunction(func):
async def async_inner(*args: P.args, **kwargs: P.kwargs) -> R:
with my_context():
result = await cast(Awaitable[R], func(*args, **kwargs))
return result
return wraps(func)(async_inner)
def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
with my_context():
result = func(*args, **kwargs)
return result
return wraps(func)(sync_inner)
Apply it like this:
@decorator
def foo(x: str) -> str: ...
@decorator
async def bar(y: int) -> int: ...
This passes mypy --strict and revealing the types of foo and bar after decoration shows what we would expect, i.e. def (x: builtins.str) -> builtins.str and def (y: builtins.int) -> typing.Awaitable[builtins.int] respectively.
The problem is that no matter what type guards we apply outside of the inner wrapper, they don't carry over to the inside of the wrapper. This has been a long-standing issue of mypy, which is not trivial to deal with.
In our scenario this means any inspect.iscoroutinefunction check done in the decorator but outside of the inner wrapper will only narrow the type to something awaitable in the scope of the decorator, but is ignored inside the wrapper. (The reason is that assignment to the non-local variable holding the function reference is possible, even after the wrapper definition. See the issue thread for details/examples.)
The typing.cast is the most straight-forward workaround in my opinion. The typing.overloads are a way to deal with the distinction between coroutines and normal functions on the caller's side. But I am curious to see, what other can come up with.