如何为Python装饰器标注可调用类型并实现全阶段类型校验

问题描述 #

需要为返回可调用对象的装饰器添加类型标注，要求：

• 装饰阶段校验被装饰函数的参数、返回值类型
• 调用阶段校验传入的参数类型

以下代码运行正常，但mypy会报错：

from typing import Any, Awaitable, Callable, Protocol

class MethodCallable(Protocol):
    def __call__(self, some_arg: int) -> Awaitable[str]: ...

def some_decorator() -> Callable[[MethodCallable], MethodCallable]:
    def wrapper(func: MethodCallable) -> MethodCallable:
        print("Decorator is applied.")
        return func
    return wrapper

class MyClass:
    @some_decorator()
    async def my_method(self, some_arg: int) -> str:
        return f"Result: {some_arg}"
    
# 示例调用
if __name__ == "__main__":
    obj = MyClass()
    print(obj.my_method(some_arg=1))  # 输出: Result: 1

mypy 错误信息 #

Argument 1 has incompatible type "Callable[[MyClass, int], Coroutine[Any, Any, str]]"; expected "MethodCallable"

需要实现的校验场景：

1. 调用传参错误：传入非int类型参数时mypy告警

   obj.my_method("str") # 应为int类型，mypy需提示错误
   

2. 装饰时定义错误：被装饰方法参数类型不符合要求时mypy告警

   class MyClass:
       @some_decorator()
       async def my_method(self, some_arg: str) -> str: # some_arg应定义为int，mypy需提示错误
   

有缺陷的临时修复方案 #

给Protocol添加__get__方法可解决原mypy报错，但存在明显问题：需要重复定义参数和返回值类型，且无法保证两处定义的一致性，会出现mypy通过但运行时报错的情况：

from typing import Any, Awaitable, Callable, Protocol
class MethodCallable(Protocol): 
    def __call__(self, __real_self: Any, __some_arg: int) -> Awaitable[str]: ...
    def __get__(self, instance: Any, owner: Any) -> "Callable[[int], Awaitable[str]]": ...


def some_decorator() -> Callable[[MethodCallable], MethodCallable]:
    def wrapper(func: MethodCallable) -> MethodCallable:
        print("Decorator is applied.")
        return func
    return wrapper

class MyClass:
    @some_decorator()
    async def my_method(self, some_arg: int) -> str:
        return f"Result: {some_arg}"

# 示例调用：mypy不报错，但运行时会因参数过多报错
if __name__ == "__main__":
    obj = MyClass()
    print(obj.my_method(1, "a"))

正确解决方案：泛型Protocol统一类型定义 #

通过ParamSpec和TypeVar实现泛型Protocol，抽象出方法签名，避免重复定义类型，同时确保__call__和__get__的类型一致性：

from typing import Any, Awaitable, Callable, Protocol, TypeVar, ParamSpec

# 定义参数和返回值的泛型
P = ParamSpec("P")
R = TypeVar("R", covariant=True)

class InstanceMethod(Protocol[P, R]):
    # 未绑定的方法（包含self参数）
    def __call__(self, __self: Any, *args: P.args, **kwargs: P.kwargs) -> Awaitable[R]: ...
    # 绑定后的方法（移除self参数）
    def __get__(self, instance: Any, owner: type[Any]) -> Callable[P, Awaitable[R]]: ...

# 具体方法类型：参数为(int)，返回值为str
MethodType = InstanceMethod[[int], str]

def some_decorator() -> Callable[[MethodType], MethodType]:
    def wrapper(func: MethodType) -> MethodType:
        print("Decorator is applied.")
        return func
    return wrapper

class MyClass:
    @some_decorator()
    async def my_method(self, some_arg: int) -> str:
        return f"Result: {some_arg}"

if __name__ == "__main__":
    obj = MyClass()
    # 以下两行会触发mypy告警，符合预期
    # obj.my_method("str")  # mypy: Argument 1 has incompatible type "str"; expected "int"
    # async def my_method(self, some_arg: str) -> str:  # mypy: Incompatible types in assignment
    print(await obj.my_method(some_arg=1))

方案优势 #

1. 避免重复定义：通过泛型统一管理参数和返回值类型，无需在__call__和__get__中重复编写
2. 类型一致性校验：泛型确保__call__（未绑定方法）和__get__（绑定方法）的类型完全匹配，不会出现定义不一致的问题
3. 双向校验：既在装饰阶段校验被装饰方法的签名，也在调用阶段校验传入的参数类型

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