如何为同步与异步函数编写类型提示装饰器？

问题描述 #

需要为同时支持同步、异步函数的装饰器添加正确的类型提示，当前实现触发mypy类型错误，希望找到合规的类型方案（而非忽略检查）。环境为Python 3.10.10，要求PyCharm能正确识别类型提示。

报错信息 #

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]

原实现代码 #

from typing import Callable, TypeVar, Awaitable, ParamSpec
import inspect
import time
from contextlib import contextmanager, wraps

P = ParamSpec("P")
R = TypeVar("R")
AR = Awaitable[R]

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)  # 报错位置
        return result

    def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
        with log_timing():
            result = foo(*args, **kwargs)
        return result  # 报错位置

    if inspect.iscoroutinefunction(foo):
        return wraps(foo)(async_inner)
    return wraps(foo)(sync_inner)

已知绕过方式（不推荐） #

通过忽略类型检查来规避错误，但不符合类型规范：

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)

正确的类型提示方案 #

使用函数重载明确区分同步、异步函数的类型签名，让类型检查器能精准匹配输入输出类型：

from typing import Callable, TypeVar, Awaitable, ParamSpec, overload
import inspect
import time
from contextlib import contextmanager, wraps

P = ParamSpec("P")
R = TypeVar("R")

@overload
def log_execution_time(foo: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
    ...

@overload
def log_execution_time(foo: Callable[P, R]) -> Callable[P, R]:
    ...

def log_execution_time(foo: Callable[P, Awaitable[R]] | Callable[P, R]) -> Callable[P, Awaitable[R]] | Callable[P, R]:
    module = inspect.getmodule(foo)
    module_spec = module.__spec__ if module else None
    module_name = module_spec.name if module_spec else foo.__module__

    @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)
        return result

    def sync_inner(*args: P.args, **kwargs: P.kwargs) -> R:
        with log_timing():
            result = foo(*args, **kwargs)
        return result

    if inspect.iscoroutinefunction(foo):
        return wraps(foo)(async_inner)
    return wraps(foo)(sync_inner)

方案说明 #

1. 用@overload定义两个重载签名：
   • 第一个针对异步函数：输入Callable[P, Awaitable[R]]，返回同类型的异步函数
   • 第二个针对同步函数：输入Callable[P, R]，返回同类型的同步函数
2. 主函数的类型提示明确输入输出的联合类型，让类型检查器能根据传入函数的类型自动匹配对应的重载逻辑
3. 此方案同时兼容mypy静态检查和PyCharm的类型推断，不会再出现类型不兼容错误

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