如何为Python类装饰器validate_connection添加正确类型提示？

修改后的完整代码 #

from functools import wraps
from typing import Any, Callable, ParamSpec, TypeVar, Union
from netmiko import ConnectHandler

# 定义泛型参数：P表示被装饰函数的参数，R表示返回值
P = ParamSpec("P")
R = TypeVar("R")

class my_class():
    def __init__(self):
        self.connection = ConnectHandler()
        # 补充原代码中用到的属性定义，避免MyPy报错
        self.name: str = ""
        self.ip: str = ""

    def validate_connection(self) -> Callable[[Callable[P, R]], Callable[P, R]]:
        def decorator(func: Callable[P, R]) -> Callable[P, R]:
            @wraps(func)
            def inner(self: my_class, *args: P.args, **kwargs: P.kwargs) -> R:
                if self.connection is None:
                    raise Exception(
                        f"No connection to {self.name}({self.ip}):22\n"
                    )
                return func(self, *args, **kwargs)
            return inner
        return decorator

    @validate_connection()
    def send_command(self, *args: Any, **kwargs: Any) -> Union[str, list[Any], dict[str, Any]]:
        return self.connection.send_command(*args, **kwargs, read_timeout=60)

关键修改点与推理过程 #

1. 补全泛型类型定义
   使用ParamSpec(P)和TypeVar(R)捕获被装饰函数的参数签名和返回值类型，让装饰器可以适配任意符合要求的实例方法，避免硬编码类型导致的兼容性问题。

2. 修正装饰器方法的定义

   • 原validate_connection缺少self参数，作为类内实例方法，这不符合Python方法规范，MyPy会错误地将传入的被装饰函数当成self参数，从而抛出类型不匹配错误。因此必须添加self: my_class参数。
   • 给validate_connection标注返回类型：Callable[[Callable[P, R]], Callable[P, R]]，明确它是一个装饰器工厂——接受一个函数，返回一个同签名的函数。

3. 调整装饰器的调用方式
   原代码中直接使用@validate_connection，现在因为validate_connection是返回装饰器的工厂方法，需要改为@validate_connection()，才能正确生成装饰器函数。

4. 内部函数的类型标注
   给inner函数的self标注my_class类型，让MyPy能正确识别self.connection、self.name等属性的类型；同时用*args: P.args和**kwargs: P.kwargs匹配被装饰函数的参数，-> R匹配返回值类型。

5. 补充类属性定义
   原代码中使用了self.name和self.ip但未初始化，添加self.name: str = ""和self.ip: str = ""避免MyPy的"未定义属性"报错。

内容的提问来源于stack exchange，提问作者Net-worker