Python 3.10：如何为支持同步/异步的重载装饰器添加类型提示？

环境 #

Python 3.10

背景 #

我已编写两个日志装饰器，分别适配普通同步函数和异步函数，Mypy校验无问题：

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

Param = ParamSpec("Param")
RetType = TypeVar("RetType")


def _log_with_bound_arguments(
    func_name: str, func_sig: Signature, *args: Any, **kwargs: Any
) -> None:
    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

需求 #

我希望将这两个装饰器合并为一个可同时用于同步和异步函数的重载装饰器ctrace，编写的代码如下：

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

问题 #

现在Mypy报出四个无法解决的错误：

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]

其中第109行是重载函数ctrace的定义行，第118行是return await func(*args, **kwargs)，第127行是return func(*args, **kwargs)。

Pylance仅显示一个错误，位于第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"

解决方案 #

问题核心是Mypy无法通过iscoroutinefunction自动窄化函数类型，且单个RetType变量同时覆盖同步/异步返回场景导致类型混淆。以下是修正后的代码：

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

Param = ParamSpec("Param")
RetType = TypeVar("RetType")
# 新增独立的异步返回类型变量，避免类型混淆
AsyncRetType = TypeVar("AsyncRetType")


def _log_with_bound_arguments(
    func_name: str, func_sig: Signature, *args: Any, **kwargs: Any
) -> None:
    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)

# 定义类型守卫，帮助Mypy识别异步函数类型
def is_async_func(func: Callable[Param, Any]) -> TypeGuard[Callable[Param, Awaitable[Any]]]:
    return iscoroutinefunction(func)


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


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


def ctrace(
    func: Union[Callable[Param, RetType], Callable[Param, Awaitable[AsyncRetType]]],
) -> Union[Callable[Param, RetType], Callable[Param, Awaitable[AsyncRetType]]]:
    if is_async_func(func):
        @wraps(func)
        async def async_wrapper(*args: Param.args, **kwargs: Param.kwargs) -> AsyncRetType:
            _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

修正要点： #

1. 拆分类型变量：新增AsyncRetType单独处理异步函数的返回类型，避免单个RetType同时对应同步返回值和异步协程的类型冲突。
2. 添加类型守卫：用TypeGuard定义is_async_func函数，让Mypy能在if分支中准确推断func的异步类型，解决类型窄化问题。
3. 匹配重载类型：调整实现函数的参数和返回类型，与重载签名严格对应，消除Mypy的类型不匹配报错。

内容的提问来源于stack exchange，提问作者kthy