如何在不修改ASTNode类的前提下为AST添加类型检查信息？

针对无法/不愿修改ASTNode基类及子类，又要给语法树节点添加类型检查数据的需求，以下是几种Pythonic的优雅方案，均支持有效的类型提示：

方案1：WeakKeyDictionary 存储类型映射（推荐） #

利用weakref.WeakKeyDictionary维护AST节点与类型信息的关联，既不用修改原类，也不会因强引用导致内存泄漏，同时类型提示清晰。

from typing import Dict, Optional, TypeVar
from weakref import WeakKeyDictionary
from your_ast_module import ASTNode

# 定义类型变量，适配你的实际类型信息（比如自定义Type类）
TypeInfo = TypeVar("TypeInfo")
ASTNodeT = TypeVar("ASTNodeT", bound=ASTNode)

# 全局映射表，键为ASTNode弱引用，值为类型信息
_node_type_map: WeakKeyDictionary[ASTNode, TypeInfo] = WeakKeyDictionary()

def get_node_type(node: ASTNodeT) -> Optional[TypeInfo]:
    """获取指定AST节点的类型信息"""
    return _node_type_map.get(node)

def set_node_type(node: ASTNodeT, type_info: TypeInfo) -> None:
    """给指定AST节点设置类型信息"""
    _node_type_map[node] = type_info

使用方式 #

类型检查阶段遍历语法树时，调用set_node_type为每个节点绑定类型信息；后续需要获取类型时，调用get_node_type即可：

# 类型检查过程
for node in traverse_ast(root_node):
    type_info = infer_type(node)  # 你的类型推断逻辑
    set_node_type(node, type_info)

# 后续获取类型
if lit_node := find_literal_node(root_node):
    if type_info := get_node_type(lit_node):
        print(f"Literal type: {type_info}")

方案2：猴子补丁+类型提示扩展（适合无__slots__的AST类） #

如果原ASTNode类没有定义__slots__，可以通过运行时动态添加属性的方式，直接在原节点上挂载类型信息，同时通过类型提示文件扩展原类的类型定义。

步骤1：添加类型提示扩展 #

创建一个类型提示文件（比如ast_ext.pyi），扩展ASTNode的类型定义：

# ast_ext.pyi
from your_ast_module import ASTNode
from typing import Optional, TypeVar

TypeInfo = TypeVar("TypeInfo")

class ASTNode(ASTNode):
    type_info: Optional[TypeInfo]  # 新增类型属性的提示

步骤2：运行时动态添加属性 #

在你的业务代码中，给ASTNode类动态初始化属性：

# 业务代码开头
from typing import Optional
from your_ast_module import ASTNode

# 给ASTNode添加默认值为None的type_info属性
ASTNode.type_info = None  # type: ignore[attr-defined]

使用方式 #

直接访问节点的type_info属性即可：

# 类型检查阶段
node.type_info = infer_type(node)

# 获取类型
if node.type_info:
    process_type(node.type_info)

注意：如果原ASTNode类定义了__slots__，此方案无法动态添加属性，此时方案1更稳妥。

方案3：带反向映射的包装类（适合需扩展节点方法的场景） #

如果需要给节点扩展额外方法（而非仅存储数据），可以优化方案B的思路，用WeakKeyDictionary维护原节点到包装类的反向映射，解决反向查找问题：

from typing import Optional, TypeVar, Generic
from weakref import WeakKeyDictionary
from your_ast_module import ASTNode

ASTNodeT = TypeVar("ASTNodeT", bound=ASTNode)
TypeInfo = TypeVar("TypeInfo")

class TypedASTNode(Generic[ASTNodeT]):
    # 全局反向映射：原ASTNode -> TypedASTNode实例
    _node_to_typed: WeakKeyDictionary[ASTNode, "TypedASTNode"] = WeakKeyDictionary()

    def __init__(self, node: ASTNodeT, type_info: TypeInfo):
        self.node = node
        self.type_info = type_info
        # 注册反向映射
        TypedASTNode._node_to_typed[node] = self

    @classmethod
    def get_typed_node(cls, node: ASTNodeT) -> Optional["TypedASTNode[ASTNodeT]"]:
        """从原ASTNode获取对应的TypedASTNode"""
        return cls._node_to_typed.get(node)

使用方式 #

# 创建包装节点
typed_node = TypedASTNode(original_node, inferred_type)

# 反向查找：从原节点获取包装类实例
if typed_node := TypedASTNode.get_typed_node(original_node):
    print(typed_node.type_info)
    # 调用包装类扩展的方法
    typed_node.validate_type()

方案对比 #

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