How to type hint an overloaded decorator that may wrap sync or async function

(Cross-posting from SO, https://stackoverflow.com/q/78206137/606576)

Python 3.10.

I have written two versions of a logging decorator, one for normal functions and one for async functions. Mypy is happy with these:

from functools import wraps
from inspect import Signature, signature
from logging import getLogger
from typing import Any, Awaitable, Callable, ParamSpec, TypeVar, Union, overload

# You heard it here first: type hinting for decorators is a pain in the arse.
Param = ParamSpec("Param")
RetType = TypeVar("RetType")


def _log_with_bound_arguments(
    func_name: str, func_sig: Signature, *args: Any, **kwargs: Any
) -> None:
    # Cf. <https://docs.python.org/3/library/inspect.html#inspect.BoundArguments>
    bound_func = func_sig.bind_partial(*args, **kwargs)
    func_params = ", ".join([k + "=" + repr(v) for k, v in bound_func.arguments.items()])
    getLogger().debug("-> %s(%s)", func_name, func_params)


def atrace(func: Callable[Param, Awaitable[RetType]]) -> Callable[Param, Awaitable[RetType]]:
    """Async decorator that safely logs the function call at the debug level."""

    @wraps(func)
    async def async_wrapper(*args: Param.args, **kwargs: Param.kwargs) -> RetType:
        _log_with_bound_arguments(func.__qualname__, signature(func), *args, **kwargs)
        return await func(*args, **kwargs)

    return async_wrapper


def trace(func: Callable[Param, RetType]) -> Callable[Param, RetType]:
    """Decorator that safely logs the function call at the debug level."""

    @wraps(func)
    def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> RetType:
        _log_with_bound_arguments(func.__qualname__, signature(func), *args, **kwargs)
        return func(*args, **kwargs)

    return wrapper

I would like to consolidate these into a single decorator that I can slap on all functions, async or not. Continuing:

from inspect import iscoroutinefunction
from typing import Union, overload


@overload
def ctrace(func: Callable[Param, RetType]) -> Callable[Param, RetType]: ...


@overload
def ctrace(func: Callable[Param, Awaitable[RetType]]) -> Callable[Param, Awaitable[RetType]]: ...


def ctrace(
    func: Union[Callable[Param, RetType], Callable[Param, Awaitable[RetType]]],
) -> Union[Callable[Param, RetType], Callable[Param, Awaitable[RetType]]]:
    if iscoroutinefunction(func):

        @wraps(func)
        async def async_wrapper(*args: Param.args, **kwargs: Param.kwargs) -> RetType:
            _log_with_bound_arguments(func.__qualname__, signature(func), *args, **kwargs)
            return await func(*args, **kwargs)

        return async_wrapper
    else:

        @wraps(func)
        def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> RetType:
            _log_with_bound_arguments(func.__qualname__, signature(func), *args, **kwargs)
            return func(*args, **kwargs)

        return wrapper

However, now Mypy is giving me four errors I can’t figure out how to fix:

src/my_package/log/__init__.py:109: error: Overloaded function implementation does not accept all possible arguments of signature 2  [misc]
src/my_package/log/__init__.py:109: error: Overloaded function implementation cannot produce return type of signature 2  [misc]
src/my_package/log/__init__.py:118: error: Returning Any from function declared to return "RetType"  [no-any-return]
src/my_package/log/__init__.py:127: error: Incompatible return value type (got "RetType | Awaitable[RetType]", expected "RetType")  [return-value]

The line numbers obviously don’t help a lot, so I’ll point out that line 109 is the def ctrace( line that starts the implementation of the overloaded function. 118 is return await func(*args, **kwargs) and 127 is return func(*args, **kwargs).

Pylance only shows a single error: on line 127, Expression of type “RetType@ctrace | Awaitable[RetType@ctrace]” cannot be assigned to return type “RetType@ctrace”
Type “RetType@ctrace | Awaitable[RetType@ctrace]” cannot be assigned to type “RetType@ctrace”
This is the same as the fourth error from Mypy as far as I can tell.

Any help in hinting this properly would be greatly appreciated.

  • mypy seems to be confused, I think because it assigns RetType=Awaitable[...]. Don’t think there is currently a good solution to this.

  • pylance is correct with it’s error message because of the way TypeGuard (i.e. iscoroutinefunction) works, see PEP 724. Essentially, the type checkers currently can’t know that func isn’t a coroutine in the else branch.

I think you wont get around adding a cast and/or type ignore somewhere.

2 Likes

Also worth taking note of PEP 742, which recently replaced PEP 724.

In fact, PEP 742 has been making a lot of recent progress. Just last week the TC voted in favor of PEP 742, and experimental support for PEP 742 has already been merged into typing_extensions (v4.10.0+), pyright (v1.1.351+), and mypy (next release, 1.10.0?).

As long as you’re OK with using experimental typing features, it may be possible to make this work using TypeIs and mypy master.

1 Like

Yep, right, should have used 742 instead, I just used the second google link for “PEP typeguard”. But using it as a solution wasn’t actually the point, it was just about explaining the problem (which both motivation sections cover very well).