带self引用的类装饰器Python类型提示问题（mypy报错）

需求背景 #

我想要实现一个可以轻松记录函数（尤其是类方法）调用次数的系统。最初的实现可以正常运行，但在修改装饰器以支持自定义方法名称后，mypy抛出了类型错误，需要修复装饰器的类型提示。

初始可正常运行的实现 #

最初编写了FunctionCount类，通过静态方法count作为装饰器记录子类方法的调用：

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

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

对应的子类使用示例，运行正常且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
    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}

支持自定义方法名称后的问题 #

为了支持自定义记录的方法名称，将count修改为带参数的装饰器：

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

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(
        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抛出错误：

Invalid self argument "A" to attribute function "func" with type "Callable[[Never], None]" mypy(misc)

解决方案 #

错误原因 #

原来的泛型变量Param和RetType是全局声明的，在带参数的装饰器结构中，mypy无法将它们正确绑定到被装饰方法的参数和返回类型，导致推断出Never这种无效类型。需要将这两个泛型变量提升到count方法的泛型参数列表中，让装饰器在接收函数时能正确捕获类型信息。

修复后的代码 #

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

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[Param: ParamSpec("Param"), RetType](
        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

关键修改说明 #

1. 将Param和RetType从全局声明移到count方法的泛型参数中：def count[Param: ParamSpec("Param"), RetType](...)（Python 3.12+支持这种泛型方法语法，若使用更低版本，可通过Generic结合TypeVar的方式兼容）。
2. 这样count方法返回的装饰器就能正确关联被装饰方法的参数规格和返回类型，mypy可以准确推断self的类型，不再抛出错误。

验证修复后的子类使用 #

修改后的装饰器可以同时支持默认名称和自定义名称的调用：

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

    @FunctionCount.count  # 使用默认名称：A.func
    def func(self) -> None:
        pass

    @FunctionCount.count(f_name="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)
assert count_dict == {"A.func": 2, "custom_name": 1}

此时运行正常，mypy也不会再抛出类型错误。

内容的提问来源于stack exchange，提问作者Paulo-99