Python类装饰器动态修改方法时的类型提示：Mypy插件实现精确静态分析

本教程深入探讨了在Python中，如何为那些在运行时动态修改类结构（如移除现有方法、添加新方法）的装饰器提供准确的类型提示。由于标准类型注解机制无法直接表达属性的删除操作，文章详细介绍了如何通过开发Mypy插件来介入静态分析过程，从而实现对装饰器行为的精确建模和类型验证，确保代码的类型安全和可维护性。

理解挑战：类装饰器与类型提示的局限性

在Python中，类装饰器是一种强大的元编程工具，它允许我们在类定义时修改或增强类的行为。常见的操作包括添加新方法、修改现有方法或移除方法。然而，当装饰器执行诸如delattr这样的动态操作时，标准的类型提示机制（包括typing模块提供的各种工具，甚至是理论上的“交叉类型”）往往难以准确地表达这些运行时行为。

考虑以下场景：一个类装饰器被设计用来移除类中的do_check方法，并添加一个基于do_check逻辑的do_assert方法。

import typing_extensions as t

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

_T = t.TypeVar("_T")

def decorator(clazz: type[_T]) -> type[_T]:
    # 运行时获取 do_check 方法
    do_check: t.Callable[[_T], bool] = getattr(clazz, "do_check")

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

    # 移除原始的 do_check 方法
    delattr(clazz, "do_check")
    # 添加新的 do_assert 方法
    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()  # 在运行时可以正常调用，但类型检查器可能缺乏提示

在上述代码中，尽管delattr(clazz, "do_check")在运行时移除了MyClass实例上的do_check方法（实际会暴露MyProtocol中的抽象方法），并添加了do_assert，但静态类型检查器（如Mypy）在没有额外信息的情况下，无法感知到这种动态变化。它可能仍然认为mc.do_check()是合法的，而对mc.do_assert()则无法提供正确的类型提示。这是因为Python的类型系统主要关注编译时（或定义时）的结构，而动态修改超出了其表达能力。

Mypy插件：高级类型分析的解决方案

为了解决上述问题，我们需要一种机制来扩展静态类型检查器的能力，使其能够理解和模拟装饰器的运行时行为。Mypy插件正是为此而生。Mypy插件允许开发者介入Mypy的类型分析过程，通过自定义钩子（hooks）来修改或增强Mypy对代码的理解。对于类装饰器动态修改类结构的情况，Mypy插件能够：

识别特定的类装饰器。
在装饰器被应用到类定义后，修改Mypy内部对该类结构的表示。
准确地添加或删除Mypy对类成员的认知，从而提供精确的类型提示。

Mypy插件实现详解

下面我们将通过一个具体的Mypy插件实现，来解决前面提到的类装饰器类型提示问题。

1. 项目结构

为了组织Mypy插件和相关代码，建议采用以下目录结构：

project/
  mypy.ini
  mypy_plugin.py
  test.py
  package/
    __init__.py
    decorator_module.py

2. mypy.ini 配置

mypy.ini 文件用于配置Mypy，告诉它加载我们的插件。

# project/mypy.ini
[mypy]
plugins = mypy_plugin.py

3. package/decorator_module.py：装饰器源码

这个文件包含了我们之前定义的MyProtocol和decorator。请注意，这里的decorator函数本身的类型注解（type[_T] -> type[_T]）在Mypy插件生效时，主要作为运行时参考，Mypy插件将接管其静态分析行为。

# project/package/decorator_module.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

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

4. mypy_plugin.py：核心插件逻辑

这是实现Mypy插件的关键文件。它定义了Mypy如何识别我们的装饰器，并在类型检查过程中修改类的结构。

# project/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

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:
        """
        这个钩子用于处理类装饰器。我们选择 `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` 时调用的钩子函数。
    它负责修改 Mypy 对类的内部表示。
    """
    # 1. 添加新的方法 `do_assert`
    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. 从 Mypy 的类定义中移除 `do_check`
    # ctx.cls.info.names 是 Mypy 存储类成员信息的地方
    del ctx.cls.info.names["do_check"]  

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

插件逻辑解析：

plugin(version: str): 这是Mypy插件的入口点，它返回我们的DecoratorPlugin类。
get_class_decorator_hook_2(self, fullname: str): Mypy提供了多种钩子来扩展其行为。对于类装饰器，get_class_decorator_hook_2是一个合适的选择，因为它在Mypy完成对类体内部的初步语义分析后被调用。我们通过检查fullname来确保只处理我们目标装饰器（package.decorator_module.decorator）。
class_decorator_hook(ctx: mypy.plugin.ClassDefContext):
- ctx对象包含了当前类定义的所有上下文信息。
- mypy.plugins.common.add_method_to_class(...): 这是一个Mypy提供的实用函数，用于向类的Mypy内部表示中添加一个新方法。我们在这里定义了do_assert的名称、参数（除了self之外没有）、返回类型以及self的类型。
- del ctx.cls.info.names["do_check"]: 这是最关键的一步。ctx.cls.info.names是一个字典，Mypy用它来存储类中所有成员的名称和对应的Mypy节点。通过直接删除"do_check"这个键，我们告诉Mypy，这个方法在装饰器处理后就不再存在于当前类中。
- 关于MyProtocol和delattr的额外说明: 当我们使用delattr(clazz, "do_check")从MyClass中删除do_check时，如果MyClass继承自MyProtocol，那么MyClass实际上会“重新暴露”MyProtocol中定义的抽象方法do_check。由于MyProtocol中的do_check是一个抽象方法（因为它raise NotImplementedError），Mypy插件通过删除MyClass自己的do_check实现，将导致Mypy认为MyClass没有实现MyProtocol的do_check，从而使其成为一个抽象类。这意味着尝试实例化MyClass将导致Mypy报错，这与运行时行为（尝试调用mc.do_check()会触发NotImplementedError）高度一致。

5. test.py：验证代码

这个文件将使用我们的装饰器，并展示Mypy在应用插件后的行为。

# project/test.py
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()

运行与验证：Mypy的精确诊断

现在，我们可以在project目录下运行Mypy来验证插件是否按预期工作：

cd project
mypy test.py

Mypy的输出将会是这样的：

test.py:7: error: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check"  [abstract]
Found 1 error in 1 file (checked 1 source file)

诊断结果解析：

test.py:7: error: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check" [abstract]:
- 这表明Mypy插件成功地识别到@decorator移除了MyClass中对MyProtocol的do_check方法的实现。
- 由于MyProtocol的do_check是抽象的，MyClass在移除了自己的实现后，Mypy正确地将其标记为一个抽象类，因此不能直接实例化。这与运行时行为（尝试实例化并调用mc.do_check()会抛出NotImplementedError）完美匹配。
mc.do_check(): 如果我们注释掉mc = MyClass()这行，Mypy将不再报错。但如果你尝试调用mc.do_check()，Mypy会警告你正在调用一个抽象方法，或者如果上下文允许，它会识别出这会引发NotImplementedError。
mc.do_assert(): 经过插件处理后，Mypy现在能够正确识别MyClass实例上存在do_assert方法，并能为其提供正确的类型提示，包括参数和返回类型。