如何为兼容同步/异步函数的装饰器添加正确类型注解？

我需要给能同时作用于同步和异步函数的装饰器添加正确的类型注解。我定义了以下类型：

type AsyncFuncType[**P, R] = Callable[P, Awaitable[R]]
type SyncFuncType[**P, R] = Callable[P, R]
type FuncType[**P, R] = AsyncFuncType[P, R] | SyncFuncType[P, R]

装饰器的包装器实际返回的是AsyncFuncType或SyncFuncType其中一种，但被要求返回FuncType，导致MyPy抛出类型错误。

原代码 #

import asyncio
from functools import wraps
from typing import Any, Awaitable, Callable, Type


type AsyncFuncType[**P, R] = Callable[P, Awaitable[R]]
type SyncFuncType[**P, R] = Callable[P, R]
type FuncType[**P, R] = AsyncFuncType[P, R] | SyncFuncType[P, R]

def exc_wrapper[**P, R](
    exc_in: Type[Exception], exc_out: Type[Exception], *exc_args: Any, **exc_kwargs: Any
) -> Callable[[FuncType[P, R]], FuncType[P, R]]:
    def decorator(
        func: FuncType[P, R],
    ) -> FuncType[P, R]:
        @wraps(func)
        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            try:
                return func(*args, **kwargs)
            except exc_in:
                raise exc_out(*exc_args, **exc_kwargs)

        @wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Awaitable[R]:
            try:
                return await func(*args, **kwargs)
            except exc_in:
                raise exc_out(*exc_args, **exc_kwargs)

        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper

    return decorator

MyPy报错信息 #

src\core\utils\decorators.py:19: error: Incompatible return value type (got "Awaitable[R] | R", expected "R")  [return-value]
src\core\utils\decorators.py:26: error: Incompatible types in "await" (actual type "Awaitable[R] | R", expected type "Awaitable[Any]")  [misc]
src\core\utils\decorators.py:31: error: Incompatible return value type (got "_Wrapped[P, object, P, Coroutine[Any, Any, Awaitable[R]]]", expected "Callable[P, Awaitable[R]] | Callable[P, R]")  [return-value]

解决方案 #

错误原因分析 #

1. sync_wrapper中，func是FuncType联合类型，MyPy无法确定调用它会返回R还是Awaitable[R]，因此报错返回类型不匹配。
2. async_wrapper的返回类型标注错误：异步函数的类型注解应该写最终返回的结果类型R，而非Awaitable[R]（异步函数本身会自动包装成Awaitable）。
3. MyPy无法通过asyncio.iscoroutinefunction的运行时判断来静态推断返回的包装器类型，需要明确类型提示。

修正后的代码 #

import asyncio
from functools import wraps
from typing import Any, Awaitable, Callable, Type, cast


type AsyncFuncType[**P, R] = Callable[P, Awaitable[R]]
type SyncFuncType[**P, R] = Callable[P, R]
type FuncType[**P, R] = AsyncFuncType[P, R] | SyncFuncType[P, R]

def exc_wrapper[**P, R](
    exc_in: Type[Exception], exc_out: Type[Exception], *exc_args: Any, **exc_kwargs: Any
) -> Callable[[FuncType[P, R]], FuncType[P, R]]:
    def decorator(func: FuncType[P, R]) -> FuncType[P, R]:
        @wraps(func)
        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            # 用cast告诉MyPy，此处func是同步函数
            sync_func = cast(SyncFuncType[P, R], func)
            try:
                return sync_func(*args, **kwargs)
            except exc_in:
                raise exc_out(*exc_args, **exc_kwargs)

        @wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            # 用cast告诉MyPy，此处func是异步函数
            async_func = cast(AsyncFuncType[P, R], func)
            try:
                return await async_func(*args, **kwargs)
            except exc_in:
                raise exc_out(*exc_args, **exc_kwargs)

        if asyncio.iscoroutinefunction(func):
            # 明确返回异步包装器类型
            return cast(AsyncFuncType[P, R], async_wrapper)
        else:
            # 明确返回同步包装器类型
            return cast(SyncFuncType[P, R], sync_wrapper)

    return decorator

关键修复点 #

• 在sync_wrapper和async_wrapper内部，用cast将func转换为对应的同步/异步函数类型，让MyPy明确调用后的返回值类型。
• 修正async_wrapper的返回类型为R，符合异步函数的类型注解规范。
• 在返回包装器时，用cast明确告诉MyPy返回的是对应类型，解决联合类型的推断问题。

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