Python修饰器类型提示

发布于02月09日

请考虑以下代码:

from functools import wraps
from netmiko import ConnectHandler

class my_class():

    def __init__(self):
        self.connection = ConnectHandler()

    def validate_connection():
        @wraps(func)
        def inner(self, *args, **kwargs):
            if self.connection is None:
                raise Exception(
                    f"No connection to {self.name}({self.ip}):22\n"
                )
            return func(self, *args, **kwargs)

        return inner


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

正确的打字方式是什么？提示validate_connection()方法，并请解释你是如何计算出来的.

目前，MyPy在第@validate_connection行给出了以下错误

Argument 1 to "validate_connection" has incompatible type "Callable[[my_class, VarArg(Any), KwArg(Any)], str | list[Any] | dict[str, Any]]"; expected "my_class"

推荐答案

我会把decorator 拉到一个顶级函数中，然后添加一个显式的typehint到self.然后，您可以为函数参数创建一个typvar，这将确保类上的结果方法保持其原始签名.

这个示例中有一个强制转换，它确实假设它将始终用于特定的类(假设self是"MyClass")，但它确实极大地简化了类型提示.

from functools import wraps
from typing import Any, TypeVar, Callable, cast, List

Func = TypeVar("Func", bound=Callable[..., Any])


def validate_connection(func: Func) -> Func:
    @wraps(func)
    def inner(self: "MyClass", *args: Any, **kwargs: Any) -> Any:
        if self.connection is None:
            raise Exception
        return func(self, *args, **kwargs)

    return cast(Func, inner)


class MyClass:
    def __init__(self) -> None:
        self.connection = None

    @validate_connection
    def send_command(self, foo: int, bar: str) -> List[str]:
        return []

使用mypy输出验证的示例:

MyClass().send_command(1, "2").append("3")
# No error :)

MyClass().send_command(1, "2").append(3)
# error: Argument 1 to "append" of "list" has incompatible type "int"; expected "str"  [arg-type]

MyClass().send_command(1, 2).append("3")
# error: Argument 2 to "send_command" of "MyClass" has incompatible type "int"; expected "str"  [arg-type]

请注意，您也可以 Select 将此设置为@staticmethod，其余所有设置仍然适用

class MyClass:
    def __init__(self) -> None:
        self.connection = None

    @staticmethod
    def validate_connection(func: Func) -> Func:
        @wraps(func)
        def inner(self: "MyClass", *args: Any, **kwargs: Any) -> Any:
            if self.connection is None:
                raise Exception
            return func(self, *args, **kwargs)

        return cast(Func, inner)

    @validate_connection
    def send_command(self, foo: int, bar: str) -> List[str]:
        return []