带有自引用的类装饰器的 Python 类型提示

Question

我的最终目标是编写一个系统来轻松记录函数调用（特别是类方法）。

我首先编写一个带有包装方法的类

Loggable

，它允许我装饰子类方法并记录它们的调用

Param = ParamSpec("Param")
RetType = TypeVar("RetType")


CountType = TypeVar("CountType", bound="FunctionCount")


class FunctionCount(Generic[CountType]):
    def __init__(self, count_dict: dict[str, int]) -> None:
        self.count_dict = count_dict

    @staticmethod
    def count(
        func: Callable[Concatenate[CountType, Param], RetType],
    ) -> Callable[Concatenate[CountType, Param], RetType]:
        def wrapper(
            self: CountType, *args: Param.args, **kwargs: Param.kwargs
        ) -> RetType:
            function_name = f"{self.__class__.__name__}.{func.__name__}"
            if function_name not in self.count_dict:
                self.count_dict[function_name] = 0
            self.count_dict[function_name] += 1

            return func(self, *args, **kwargs)

        return wrapper

现在我可以编写子类并记录它们的调用：

class A(FunctionCount):
    def __init__(self, count_dict: dict[str, int]) -> None:
        super().__init__(count_dict)

    @FunctionCount.count
    def func(self) -> None:
        pass

    @FunctionCount.count
    def func2(self) -> None:
        pass


count_dict: dict[str, int] = {}

a = A(count_dict)

a.func()
a.func()
a.func2()

print(count_dict)
assert count_dict == {"A.func": 2, "A.func2": 1}

效果非常好，我很高兴。 但后来我认为为方法提供自定义名称会很好，所以我将包装器更改为装饰器

class FunctionCount(Generic[CountType]):
    def __init__(self, count_dict: dict[str, int]) -> None:
        self.count_dict = count_dict

    @staticmethod
    def count(
        f_name: str | None = None,
    ) -> Callable[
        [Callable[Concatenate[CountType, Param], RetType]],
        Callable[Concatenate[CountType, Param], RetType],
    ]:
        def decorator(
            func: Callable[Concatenate[CountType, Param], RetType],
        ) -> Callable[Concatenate[CountType, Param], RetType]:
            def wrapper(
                self: CountType, *args: Param.args, **kwargs: Param.kwargs
            ) -> RetType:
                function_name = f_name or f"{self.__class__.__name__}.{func.__name__}"
                if function_name not in self.count_dict:
                    self.count_dict[function_name] = 0
                self.count_dict[function_name] += 1

                return func(self, *args, **kwargs)

            return wrapper

        return decorator

然后我只需更改装饰器调用

class A(FunctionCount):
    def __init__(self, count_dict: dict[str, int]) -> None:
        super().__init__(count_dict)

    @FunctionCount.count()
    def func(self) -> None:
        pass

    @FunctionCount.count("custom_name")
    def func2(self) -> None:
        pass

a.func()
a.func()
a.func2()

print(count_dict)
assert count_dict == {"A.func": 2, "custom_name": 1}

这个脚本也运行得很好，但现在 mypy 给我带来了困难。 当我调用

a.func

方法时，出现以下 mypy 错误：

类型为“Callable[[Never], None]”的属性函数“func”的自参数“A”无效 mypy(misc)

我猜使用装饰器而不是包装器导致了这个错误，但我不明白为什么以及我应该做什么来纠正它。

有人知道如何以这种方式拥有一个正确类型的装饰器吗？

Answer 1

mypy

有问题，因为类型变量和参数规范仅出现在返回注释中（而不出现在参数中）。不幸的是，我不知道如何让它适用于每种边缘情况，特别是当您用它注释（实例或类）方法时。但对于常规函数（至少它适用于我的大多数用例），可能可以创建一个通用装饰器协议并将其设置为装饰器的返回类型 - 类似于：

class Decorator(Protocol, Generic[CountType, Param, RetType]):
    __call__(self, __self__: CountType, *args: Param.args, **kwargs: Param.kwargs) -> RetType:
        ...


class FunctionCount(Generic[CountType]):
    def __init__(self, count_dict: dict[str, int]) -> None:
        self.count_dict = count_dict

    @staticmethod
    def count(
        f_name: str | None = None,
    ) -> Decorator[CountType, Param, RetType]:
        def decorator(
            func: Callable[Concatenate[CountType, Param], RetType],
        ) -> Callable[Concatenate[CountType, Param], RetType]:
            def wrapper(
                self: CountType, *args: Param.args, **kwargs: Param.kwargs
            ) -> RetType:
                function_name = f_name or f"{self.__class__.__name__}.{func.__name__}"
                if function_name not in self.count_dict:
                    self.count_dict[function_name] = 0
                self.count_dict[function_name] += 1

                return func(self, *args, **kwargs)

            return wrapper

        return decorator