使用 Mypy 插件为动态修改类方法的装饰器提供准确类型提示

from typing import Callable, Protocol, TypeVar

_T = TypeVar("_T")

class MyProtocol(Protocol):
    def do_check(self) -> bool:
        raise NotImplementedError

def decorator(clazz: type[_T]) -> type[_T]:
    # 在运行时获取并移除 do_check
    do_check: Callable[[_T], bool] = getattr(clazz, "do_check")

    def do_assert(self: _T) -> None:
        assert do_check(self)

    delattr(clazz, "do_check")
    setattr(clazz, "do_assert", do_assert)

    return clazz

@decorator
class MyClass(MyProtocol):
    def do_check(self) -> bool:
        return False

# 运行时行为
mc = MyClass()
# mc.do_check() # 运行时会报错，因为它已被移除（或暴露了基类的NotImplementedError）
mc.do_assert() # 运行时正常工作

project/
  mypy.ini           # Mypy 配置文件，用于启用插件
  mypy_plugin.py     # Mypy 插件的实现
  test.py            # 包含使用装饰器的示例代码
  package/
    __init__.py
    decorator_module.py # 包含协议和装饰器定义

[mypy]
plugins = mypy_plugin.py

from __future__ import annotations

import typing_extensions as t

import mypy.plugin
import mypy.plugins.common
import mypy.types

if t.TYPE_CHECKING:
    import collections.abc as cx
    import mypy.nodes

# Mypy 插件的入口点
def plugin(version: str) -> type[DecoratorPlugin]:
    return DecoratorPlugin

class DecoratorPlugin(mypy.plugin.Plugin):
    # 注册一个类装饰器 Hook
    # get_class_decorator_hook_2 在类体语义分析完成后触发
    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        # 检查装饰器是否是我们关注的 'package.decorator_module.decorator'
        if fullname == "package.decorator_module.decorator":
            return class_decorator_hook
        return None

def class_decorator_hook(ctx: mypy.plugin.ClassDefContext) -> bool:
    """
    当 Mypy 遇到 @decorator 装饰的类时，此 Hook 会被调用。
    我们在此处修改 Mypy 对该类的内部表示。
    """
    # 1. 添加 'do_assert' 方法的类型信息
    # ctx.api 提供了访问 Mypy 内部 API 的接口
    # ctx.cls 是当前被装饰的类的 AST 节点
    # ctx.cls.fullname 是类的完整名称
    mypy.plugins.common.add_method_to_class(
        ctx.api,
        cls=ctx.cls,
        name="do_assert",
        args=[],  # 实例方法，除了 self 以外没有其他参数
        return_type=mypy.types.NoneType(), # 返回类型为 None
        self_type=ctx.api.named_type(ctx.cls.fullname), # self 的类型是当前类
    )

    # 2. 从类的内部命名空间中移除 'do_check'
    # ctx.cls.info.names 存储了 Mypy 对类成员的内部表示
    del ctx.cls.info.names["do_check"]  # 移除 `do_check` 方法

    return True  # 返回 True 表示类已完全定义，不需要进一步的语义分析

from __future__ import annotations

import typing_extensions as t

if t.TYPE_CHECKING:
    import collections.abc as cx
    _T = t.TypeVar("_T")

class MyProtocol(t.Protocol):
    def do_check(self) -> bool:
        raise NotImplementedError

# 装饰器本身的类型提示可以保持简单，
# 因为 Mypy 插件会在 Mypy 内部处理类型逻辑。
def decorator(clazz: type[_T]) -> type[_T]:
    do_check: cx.Callable[[_T], bool] = getattr(clazz, "do_check")

    def do_assert(self: _T) -> None:
        assert do_check(self)

    delattr(clazz, "do_check")
    setattr(clazz, "do_assert", do_assert)

    return clazz

from package.decorator_module import MyProtocol, decorator

@decorator
class MyClass(MyProtocol):
    def do_check(self) -> bool:
        return False

mc = MyClass()  # mypy: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check" [abstract]
mc.do_check()   # mypy: error: "MyClass" has no attribute "do_check" [attr-without-any-type]
mc.do_assert()  # OK，Mypy 能够正确识别

test.py:7: error: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check"  [abstract]
test.py:8: error: "MyClass" has no attribute "do_check"  [attr-without-any-type]
Found 2 errors in 1 file (checked 1 source file)