使用 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.do_check() 会失败，但 mc.do_assert() 会成功
mc = MyClass()
# mc.do_check() # 运行时会抛出 NotImplementedError
# mc.do_assert() # 运行时正常工作

# 问题在于：
# 1. Mypy 仍然认为 mc.do_check() 存在，并可能给出错误的提示。
# 2. Mypy 无法识别 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

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 插件来说，它会通过 fullname 识别并接管处理。
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 __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

def plugin(version: str) -> type[DecoratorPlugin]:
    """Mypy 插件的入口函数。"""
    return DecoratorPlugin

class DecoratorPlugin(mypy.plugin.Plugin):
    """自定义 Mypy 插件类。"""

    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        """
        这个钩子在 Mypy 处理完类体后，遇到类装饰器时触发。
        我们通过装饰器的全名来识别并应用自定义逻辑。
        """
        # 检查是否是我们关注的装饰器
        if fullname == "package.decorator_module.decorator":
            return class_decorator_hook
        return None

def class_decorator_hook(ctx: mypy.plugin.ClassDefContext) -> bool:
    """
    当 Mypy 遇到 `@decorator` 时调用的实际处理函数。
    它负责修改 Mypy 对类的内部表示。
    """
    # 1. 添加 do_assert 方法
    # mypy.plugins.common.add_method_to_class 是一个便捷函数，
    # 用于向类添加方法。
    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 方法
    # 直接删除类信息中的名称条目。
    # 注意：delattr(clazz, "do_check") 实际上会暴露 MyProtocol 中的 do_check。
    # 这里删除的是 MyClass 自身实现的 do_check。
    del ctx.cls.info.names["do_check"]

    # 返回 True 表示类已完全定义，无需再次进行语义分析。
    return True

from package.decorator_module import MyProtocol, decorator

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

mc = MyClass()  # 预期 Mypy 报错
mc.do_check()   # 预期 Mypy 报错
mc.do_assert()  # 预期 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" (note: "do_check" is an abstract method in "MyProtocol") [attr-defined]