如何用Python ast库深度遍历AST树获取变量完整迭代值？

问题分析 #

用ast.walk遍历AST时采用广度优先策略，会先处理外层循环再处理内层循环，无法模拟嵌套循环的实际执行流程。比如这段代码：

for i in range(2):
    for j in range(2):
        pass

用ast.walk提取变量值会得到{'i': [0, 1], 'j': [0, 1]}，但实际执行中j的取值顺序是[0,1,0,1]，需要按执行流程遍历AST才能得到正确结果。

用ast.NodeVisitor实现深度优先+执行流程模拟 #

ast.NodeVisitor是Python AST模块提供的深度遍历工具，通过自定义子类并重写特定节点的访问方法，可以精确控制遍历顺序并模拟代码执行上下文。

自定义Visitor实现代码 #

import ast

class LoopExecutionVisitor(ast.NodeVisitor):
    def __init__(self):
        self.variable_tracker = {}
        # 栈结构跟踪当前嵌套循环的变量迭代上下文
        self.loop_context_stack = []

    def visit_For(self, node):
        # 解析循环迭代器的实际值（这里先处理range和列表字面量）
        iterable = self._resolve_iterable(node.iter)
        if not iterable:
            return self.generic_visit(node)
        
        var_name = node.target.id
        # 将当前循环的变量和迭代器推入栈
        self.loop_context_stack.append((var_name, iter(iterable)))

        # 模拟循环执行：逐个取出迭代值，遍历循环体
        for val in iterable:
            # 记录变量值
            if var_name not in self.variable_tracker:
                self.variable_tracker[var_name] = []
            self.variable_tracker[var_name].append(val)
            # 递归遍历循环体，此时上下文是当前循环变量的当前值
            self.generic_visit(node.body)
        
        # 循环结束，弹出上下文栈
        self.loop_context_stack.pop()

    def _resolve_iterable(self, node):
        # 处理range调用
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == 'range':
            args = [ast.literal_eval(arg) for arg in node.args]
            return range(*args)
        # 处理列表字面量
        elif isinstance(node, ast.List):
            return [ast.literal_eval(elem) for elem in node.elts]
        # 可扩展处理tuple、生成器等其他可迭代对象
        return None

# 测试示例
test_code = """
for i in range(2):
    for j in range(2):
        pass
"""
ast_tree = ast.parse(test_code)
visitor = LoopExecutionVisitor()
visitor.visit(ast_tree)
print(visitor.variable_tracker)  # 输出: {'i': [0, 1], 'j': [0, 1, 0, 1]}

代码解析 #

• LoopExecutionVisitor：通过loop_context_stack维护嵌套循环的变量上下文，确保内层循环能在外层循环的每个迭代下正确执行。
• visit_For：重写该方法模拟循环执行流程，先解析迭代器的所有值，再逐个取值并记录，之后递归遍历循环体，保证内层循环的变量值按实际执行顺序被记录。
• _resolve_iterable：负责解析AST中的迭代器节点为实际可迭代对象，可根据需求扩展支持更多类型（如元组、集合等）。

ast.NodeVisitor用法说明（法语术语版） #

• ast.NodeVisitor：Classe de base pour le parcours en profondeur d'un arbre AST.
• Méthodes surchargeables：Pour chaque type de nœud AST (ex: For, Assign, If), définissez une méthode visit_NomDuNoeud qui sera appelée lors de la visite de ce nœud.
• generic_visit：Méthode par défaut qui visite tous les enfants du nœud courant. Appelez-la dans vos méthodes surchargées pour continuer le parcours des sous-arbres.
• Gestion du contexte：Utilisez des attributs de classe (comme loop_context_stack) pour conserver l'état pendant le parcours, indispensable pour traiter des structures imbriquées comme les boucles.

其他扩展方案 #

• 环境变量跟踪：如果需要处理循环体内变量被修改的情况，可以扩展Visitor维护一个当前变量值的字典（类似解释器的局部环境），在遍历过程中实时更新变量值，再记录到结果中。
• 支持更多迭代器类型：在_resolve_iterable中添加对生成器表达式、内置可迭代对象（如enumerate）的解析逻辑，提升工具的通用性。

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