如何通过Mypy插件实现动态属性创建的类型检查？

背景与问题 #

我们的代码库实现了一套自定义的命令行选项定义机制，核心代码如下（opts.py）：

# opts.py
def group():
    o = OptionsGroup()
    return o, o.define

options = _SomeOptionsSingletonClass()

def define(name: str, type_: Type, default: bool, ...):
    # 在options单例对象上创建并设置属性以跟踪该选项
    ...

class OptionsGroup:
    def define(name: str, type_: Type, default: bool, ...):
        define(*args)

代码中有两种使用模式：

模式1：局部选项组 #

# some_module.py
from opts import group

options, define = group()

define("foo", type="str", ...)

# ...

print(options.foo) # 此处无法被Mypy做类型检查

模式2：全局单例选项 #

# some_module.py
from opts import options, define

define("foo", type="str", ...)

# ...

print(options.foo) # 此处无法被Mypy做类型检查

我们依赖Mypy做类型检查，但当前这套动态定义的选项无法被Mypy识别。现有的临时方案是用脚本修改代码后再跑Mypy，但这种方式无法直接检查原代码，且维护成本高。尝试过开发Mypy插件，但模式1中的define调用无法触发钩子，Analyze Type钩子也难以适配需求。需要一套可靠的方案，让Mypy能在处理define调用和options属性引用时自动合成对应类型。

解决方案：定制Mypy插件 #

通过开发Mypy插件，跟踪define调用的参数，为options的属性访问动态生成类型。以下是具体实现思路和代码示例：

核心思路 #

1. 跟踪全局选项：捕获全局define函数的调用，记录每个选项的名称和对应类型，在访问options.xxx时返回该类型。
2. 跟踪局部选项组：识别group()的赋值语句，绑定局部define方法到对应的options变量，记录该变量下的所有选项类型，在访问局部options.xxx时返回对应类型。

插件代码实现 #

from mypy.plugin import Plugin, CallExprContext, AssignExprContext, AttributeExprContext
from mypy.nodes import CallExpr, AssignExpr, TupleExpr, NameExpr, AttributeExpr, StrExpr
from mypy.types import Instance
from typing import Dict, Optional

class OptionsPlugin(Plugin):
    def __init__(self, options) -> None:
        super().__init__(options)
        # 存储全局选项：{选项名: Mypy类型实例}
        self.global_options: Dict[str, Instance] = {}
        # 存储局部选项：{局部options变量名: {选项名: Mypy类型实例}}
        self.local_options: Dict[str, Dict[str, Instance]] = {}
        # 映射局部define变量到对应的options变量
        self.local_defines_map: Dict[str, str] = {}

    def get_call_hook(self, fullname: str):
        # 处理全局define函数调用
        if fullname == "opts.define":
            def hook(ctx: CallExprContext) -> Optional[Instance]:
                # 解析选项名称和类型参数
                if len(ctx.args) >= 2 and isinstance(ctx.args[0], StrExpr) and isinstance(ctx.args[1], StrExpr):
                    opt_name = ctx.args[0].value
                    type_name = ctx.args[1].value
                    # 获取对应内置类型的Mypy实例
                    opt_type = ctx.api.named_type(f"builtins.{type_name}")
                    if isinstance(opt_type, Instance):
                        self.global_options[opt_name] = opt_type
                return None
            return hook

        # 处理OptionsGroup的define方法调用
        elif fullname == "opts.OptionsGroup.define":
            def hook(ctx: CallExprContext) -> Optional[Instance]:
                # 找到调用define的局部变量名
                if (isinstance(ctx.callee, AttributeExpr) 
                    and isinstance(ctx.callee.expr, NameExpr)):
                    define_var_name = ctx.callee.expr.name
                    opt_var_name = self.local_defines_map.get(define_var_name)
                    if not opt_var_name:
                        return None

                    # 解析选项参数
                    if len(ctx.args) >= 2 and isinstance(ctx.args[0], StrExpr) and isinstance(ctx.args[1], StrExpr):
                        opt_name = ctx.args[0].value
                        type_name = ctx.args[1].value
                        opt_type = ctx.api.named_type(f"builtins.{type_name}")
                        if isinstance(opt_type, Instance):
                            # 初始化局部选项存储
                            if opt_var_name not in self.local_options:
                                self.local_options[opt_var_name] = {}
                            self.local_options[opt_var_name][opt_name] = opt_type
                return None
            return hook

        return None

    def get_assign_hook(self):
        def hook(ctx: AssignExprContext) -> Optional[Instance]:
            # 捕获options, define = group()的赋值语句
            if (isinstance(ctx.lvalue, TupleExpr) 
                and len(ctx.lvalue.items) == 2
                and isinstance(ctx.lvalue.items[0], NameExpr)
                and isinstance(ctx.lvalue.items[1], NameExpr)):
                
                opt_var_name = ctx.lvalue.items[0].name
                define_var_name = ctx.lvalue.items[1].name
                # 记录局部define与options的映射
                self.local_defines_map[define_var_name] = opt_var_name
            return None
        return hook

    def get_attribute_hook(self, fullname: str):
        # 处理options的属性访问
        if fullname.startswith("opts._SomeOptionsSingletonClass.") or fullname.startswith("opts.OptionsGroup."):
            def hook(ctx: AttributeExprContext) -> Optional[Instance]:
                attr_name = ctx.context.name
                # 判断是全局还是局部options
                if isinstance(ctx.context.expr, NameExpr):
                    var_name = ctx.context.expr.name
                    # 全局options
                    if var_name == "options":
                        return self.global_options.get(attr_name)
                    # 局部options
                    else:
                        return self.local_options.get(var_name, {}).get(attr_name)
                return None
            return hook
        return None

def plugin(version: str):
    return OptionsPlugin

配置与使用 #

1. 将上述代码保存为opts_plugin.py。
2. 在项目的mypy.ini中添加配置：

   [mypy]
   plugins = opts_plugin.py
   

3. 正常运行mypy即可自动识别动态定义的选项类型。

扩展优化点 #

• 支持type_参数传入类型对象（而非字符串），需要解析表达式获取对应的Mypy类型。
• 验证选项默认值的类型与指定类型是否匹配，添加类型错误提示。
• 处理作用域问题，确保局部选项不会跨函数或类污染。
• 支持复杂类型（如List[str]、自定义类）的解析。

内容的提问来源于stack exchange，提问作者David Reed