为动态修改类结构的Python装饰器提供类型提示：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() # 运行时会报错，但类型检查器可能仍认为其存在
# mc.do_assert() # 运行时正常，但类型检查器可能无法识别

project/
  mypy.ini
  mypy_plugin.py
  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 插件来说不具有决定性作用，
# 插件会通过其内部逻辑处理类型信息。
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):
    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        """
        注册一个钩子，用于在类定义被装饰器处理后进行类型修改。
        选择 `get_class_decorator_hook_2` 是因为它在类体语义分析完成后执行。
        """
        if fullname == "package.decorator_module.decorator":
            return class_decorator_hook
        return None

def class_decorator_hook(ctx: mypy.plugin.ClassDefContext) -> bool:
    """
    当 Mypy 遇到 `@decorator` 时调用的实际钩子函数。
    它会模拟装饰器对类结构的修改。
    """
    # 添加 do_assert 方法
    mypy.plugins.common.add_method_to_class(
        ctx.api,
        cls=ctx.cls,
        name="do_assert",
        args=[],  # 实例方法，除了 self 外没有其他参数
        return_type=mypy.types.NoneType(),
        self_type=ctx.api.named_type(ctx.cls.fullname),
    )
    # 移除 do_check 方法
    # 注意：delattr(clazz, "do_check") 在运行时会移除实例或类上的属性。
    # 如果基类（如 Protocol）定义了该方法，移除后可能会暴露基类的抽象方法。
    # Mypy 插件通过直接操作 `ctx.cls.info.names` 来模拟这种移除。
    # 如果 MyProtocol.do_check 仍然存在，Mypy 会将其视为抽象方法。
    if "do_check" in ctx.cls.info.names:
        del ctx.cls.info.names["do_check"]

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

from package.decorator_module import MyProtocol, decorator

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

mc = MyClass()
mc.do_check()
mc.do_assert()

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