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

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

我首先编写一个带有包装方法的类 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

这个脚本也可以很好地工作，但现在 mypy 给了我困难时期。 当我调用

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 无效的自参数“A”到类型为“Callable[[Never], None]”的属性函数“func” mypy(misc)

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

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

Does someone know how to have a correctly typed decorator this way ?

遇到的 mypy 错误是由于在装饰器定义中， CountType 的类型变量没有正确地从方法定义传递到装饰器内部。以下是如何解决此问题并改进类型提示的方法：

from typing import Callable, Concatenate, ParamSpec, TypeVar, Generic, overload

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

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

    @overload
    @staticmethod
    def count(
        f_name: str,
    ) -> Callable[
        [Callable[Concatenate['FunctionCount', Param], RetType]], 
        Callable[Concatenate['FunctionCount', Param], RetType]
    ]: ...

    @overload
    @staticmethod
    def count(
    ) -> Callable[
        [Callable[Concatenate['FunctionCount', Param], RetType]], 
        Callable[Concatenate['FunctionCount', Param], RetType]
    ]: ...

    @staticmethod
    def count(
        f_name: str | None = None,
    ) -> Callable[
        [Callable[Concatenate['FunctionCount', Param], RetType]], 
        Callable[Concatenate['FunctionCount', Param], RetType]
    ]:
        def decorator(
            func: Callable[Concatenate['FunctionCount', Param], RetType],
        ) -> Callable[Concatenate['FunctionCount', Param], RetType]:
            def wrapper(
                self: 'FunctionCount', *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

count_dict: dict[str, int] = {}

a = A(count_dict)

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

print(count_dict)  # 输出: {'A.func': 2, 'custom_name': 1}

解释：

1. 移除 CountType : 我们不需要 CountType 类型变量，因为我们总是在 FunctionCount 类中使用它。
2. FunctionCount 作为第一个参数： 我们将 FunctionCount 类本身作为装饰器函数的第一个参数传递给 Concatenate ，因为我们希望跟踪该类的方法调用。
3. 使用 overload : 我们使用 @overload 装饰器来提供两种不同的类型提示，一种用于带有名稱参数的装饰器，另一种用于不带名称参数的装饰器。
4. 明确的类型提示： 我们在整个代码中提供了更清晰、更明确的类型提示，以帮助 mypy 更好地理解代码。

通过这些更改，mypy 将不再报告错误，并且的代码将具有更强大的类型安全性。