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

本文探讨了如何为通过装饰器动态添加或删除方法的Python类提供准确的类型提示。由于标准类型提示无法表达此类复杂的运行时类结构修改，MyPy插件成为解决这一挑战的强大工具。通过定制MyPy的行为，我们可以确保静态类型检查器正确识别装饰器修改后的类结构，从而提升代码的健壮性和可维护性。

1. 问题背景：类装饰器与类型提示的局限性

在Python中，类装饰器是一种强大的元编程工具，可以在类定义时修改或增强类的行为。一个常见的场景是，装饰器可能会从类中移除一个现有方法，并添加一个新方法。然而，为这种动态修改提供准确的类型提示，对于静态类型检查器（如MyPy）来说是一个挑战。

考虑以下示例代码，一个装饰器 decorator 旨在移除 do_check 方法并添加 do_assert 方法：

import typing_extensions as t
import collections.abc as cx

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_assert
    do_check: cx.Callable[[_T], bool] = getattr(clazz, "do_check")

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

    delattr(clazz, "do_check") # 移除 do_check
    setattr(clazz, "do_assert", do_assert) # 添加 do_assert

    return clazz

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

mc = MyClass()
mc.do_check()   # 运行时会报错，但MyPy可能仍提示该方法存在
mc.do_assert()  # 运行时正常工作，但MyPy可能无法提供类型提示

在这个例子中，decorator 运行时成功地修改了 MyClass。然而，如果没有特殊的处理，MyPy 会遇到以下问题：

• mc.do_check()：MyPy 可能会错误地认为该方法仍然存在，并提供其原始的类型提示，尽管在运行时它已被移除（或者更准确地说，被 MyProtocol 的抽象实现替代）。
• mc.do_assert()：MyPy 无法识别这个由装饰器动态添加的方法，因此无法提供类型提示或进行类型检查。

即使是使用交叉类型（Intersection Type），也无法表达“删除一个属性”这样的操作。标准类型提示机制的局限性在于它们主要用于描述静态的、预定义的类型结构，难以应对运行时发生的复杂结构变动。

2. 解决方案：利用 MyPy 插件扩展类型检查能力

为了解决上述问题，我们需要一种机制来告知 MyPy 装饰器对类结构所做的具体更改。MyPy 插件正是为此目的而设计的强大工具。通过编写一个 MyPy 插件，我们可以在 MyPy 进行类型检查时，介入并修改它对被装饰类的理解。

2.1 MyPy 插件的工作原理

MyPy 插件允许开发者在 MyPy 的语义分析阶段注入自定义逻辑。当 MyPy 遇到特定的装饰器、函数或类时，它可以调用插件中注册的钩子（hooks）。这些钩子可以访问和修改 MyPy 内部表示的抽象语法树（AST）或类型信息，从而实现自定义的类型检查行为。

对于类装饰器，MyPy 提供了 get_class_decorator_hook_2 这样的钩子。这个钩子在类体已经被语义分析之后，但在类定义最终确定之前被调用，这正是修改类结构信息的好时机。

2.2 实现 MyPy 插件

下面将详细介绍如何构建一个 MyPy 插件来正确处理上述类装饰器。

项目结构：

首先，设置以下文件目录结构：

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

mypy.ini 配置：

在 mypy.ini 文件中配置 MyPy 以加载我们的插件：

[mypy]
plugins = mypy_plugin.py

这行配置告诉 MyPy 在运行时加载并执行 mypy_plugin.py 文件中的插件。

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]:
    return DecoratorPlugin

class DecoratorPlugin(mypy.plugin.Plugin):
    # 注册类装饰器钩子
    # 当 MyPy 遇到 'package.decorator_module.decorator' 装饰器时，
    # 将调用 class_decorator_hook 函数
    def get_class_decorator_hook_2(
        self, fullname: str
    ) -> cx.Callable[[mypy.plugin.ClassDefContext], bool] | None:
        if fullname == "package.decorator_module.decorator":
            return class_decorator_hook
        return None

def class_decorator_hook(ctx: mypy.plugin.ClassDefContext) -> bool:
    # 1. 添加 do_assert 方法
    mypy.plugins.common.add_method_to_class(
        ctx.api,
        cls=ctx.cls,
        name="do_assert",
        args=[],  # 实例方法，不接受额外参数（self 参数由 MyPy 自动处理）
        return_type=mypy.types.NoneType(), # 返回类型为 None
        self_type=ctx.api.named_type(ctx.cls.fullname), # self 的类型是当前类
    )
    # 2. 从类的类型信息中移除 do_check 方法
    del ctx.cls.info.names["do_check"]

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

代码解析：

• plugin(version: str)：这是 MyPy 插件的入口点，它返回一个插件类的实例。
• DecoratorPlugin(mypy.plugin.Plugin)：自定义插件类，继承自 mypy.plugin.Plugin。
• get_class_decorator_hook_2(self, fullname: str)：这个钩子是专门用于类装饰器的。fullname 是装饰器的完全限定名（例如 module.decorator_name）。当 fullname 匹配到 package.decorator_module.decorator 时，我们返回 class_decorator_hook 函数。
• class_decorator_hook(ctx: mypy.plugin.ClassDefContext)：
  • ctx.api：提供了与 MyPy 核心交互的接口。
  • ctx.cls：表示当前被装饰的类。
  • mypy.plugins.common.add_method_to_class(...)：这是一个实用函数，用于向 MyPy 对类的理解中添加一个方法。我们指定了方法名 do_assert、无额外参数、返回类型为 None，以及 self 的类型为当前类。
  • del ctx.cls.info.names["do_check"]：这是关键一步。它直接从 MyPy 内部表示的类信息中移除了 do_check 方法。这意味着 MyPy 将不再认为该方法存在于被装饰的类上。
  • 返回 True 表示插件已完成对类的修改，MyPy 可以继续后续的类型检查。

package/decorator_module.py - 装饰器实现：

这个文件包含实际的 Python 装饰器代码。请注意，这里的类型提示主要是为了运行时行为，MyPy 插件将接管其类型检查行为。

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 插件来说不具有实际意义，
# 插件会在检测到 @package.decorator_module.decorator 时执行其自定义逻辑。
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

test.py - 测试代码：

这个文件用于验证 MyPy 插件是否按预期工作。

from package.decorator_module import MyProtocol, decorator

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

mc = MyClass()  # 预期 MyPy 报错：无法实例化抽象类 "MyClass"
mc.do_check()   # 预期 MyPy 报错或提示不存在，运行时会引发 NotImplementedError
mc.do_assert()  # 预期 MyPy 正常识别并提供类型提示

2.3 运行 MyPy 验证

现在，在 project 目录下运行 MyPy：

mypy test.py

你将看到类似以下的 MyPy 输出：

test.py:7: error: Cannot instantiate abstract class "MyClass" with abstract attribute "do_check"  [abstract]

输出解读：

1. Cannot instantiate abstract class "MyClass" with abstract attribute "do_check"

   • 这正是我们期望的结果！MyPy 插件通过 del ctx.cls.info.names["do_check"] 从 MyClass 的类型定义中移除了 do_check。
   • 由于 MyClass 继承自 MyProtocol，而 MyProtocol 定义了抽象方法 do_check (raise NotImplementedError)，一旦 MyClass 自己的 do_check 被插件“移除”，MyPy 就会认为 MyClass 没有实现 MyProtocol 的 do_check，从而使其成为一个抽象类。
   • 因此，尝试实例化一个抽象类 MyClass() 会导致 MyPy 报错。

2. mc.do_check()

   • 如果你注释掉 mc = MyClass() 这一行，MyPy 将不再报告实例化错误。但是，如果你尝试调用 mc.do_check()，MyPy 将会报错，因为它现在知道 MyClass 实例上没有 do_check 方法（或者它是一个抽象方法）。
   • 在运行时，由于 delattr(clazz, "do_check")，调用 mc.do_check() 将会触发 MyProtocol 中定义的 NotImplementedError。MyPy 插件的类型检查结果与运行时行为完美匹配。

3. mc.do_assert()

   • MyPy 将正确识别 mc.do_assert() 方法，并为其提供正确的类型提示，因为插件通过 add_method_to_class 明确告知了 MyPy 这个方法的存在及其签名。

3. 总结与注意事项

• MyPy 插件的强大之处： MyPy 插件提供了一种强大的机制，可以扩展 MyPy 的类型检查能力，以应对标准类型提示无法处理的复杂场景，例如运行时动态修改类结构。
• 弥合运行时与静态检查的鸿沟： 通过插件，我们可以确保 MyPy 的静态类型检查结果与 Python 代码的实际运行时行为保持一致，这对于提高代码质量和减少运行时错误至关重要。
• 学习成本： 编写 MyPy 插件需要一定的学习成本，包括理解 MyPy 的内部 API、AST 结构和插件钩子。但对于需要处理复杂元编程模式的项目来说，这种投入是值得的。
• 谨慎使用： 插件提供了极大的灵活性，但也应谨慎使用。过度或不当的插件可能会使类型检查变得复杂或难以理解。
• 装饰器中的类型提示： 在 decorator_module.py 中，decorator 函数本身的类型提示 (def decorator(clazz: type[_T]) -> type[_T]:) 尽管在插件生效时会被插件的逻辑覆盖，但在没有插件或插件未被激活的环境下，它仍然提供了一个基本的类型回退。

通过本文的教程，我们展示了如何利用 MyPy 插件为动态修改类方法的装饰器提供精确的类型提示，从而在复杂的 Python 项目中实现更严格、更可靠的静态类型检查。