如何用Python的ast库查找变量的最初赋值

要解决变量溯源的问题，核心是递归追踪变量的引用，直到找到它的初始常量赋值，而不是只获取直接赋值的变量名。以下是具体实现方案：

实现思路 #

1. 遍历AST构建符号表，记录每个变量的赋值节点；
2. 对目标变量进行递归解析：如果赋值是变量引用，就继续查该变量的赋值，直到找到常量值或无法追踪。

完整代码实现 #

import ast

class AssignmentTracker(ast.NodeVisitor):
    def __init__(self):
        # 符号表：变量名 -> 赋值节点
        self.variable_map = {}

    def visit_Assign(self, node):
        # 处理单目标赋值（如需支持多变量赋值如a=b=c，可扩展处理ast.Tuple类型的target）
        for target in node.targets:
            if isinstance(target, ast.Name):
                var_name = target.id
                self.variable_map[var_name] = node.value
        self.generic_visit(node)

def get_initial_value(variable_map, var_name):
    # 递归解析变量的初始值
    current_node = variable_map.get(var_name)
    if not current_node:
        return None
    
    # 处理字符串常量
    if isinstance(current_node, ast.Constant) and isinstance(current_node.value, str):
        return current_node.value
    # 处理变量引用，递归追踪
    elif isinstance(current_node, ast.Name):
        return get_initial_value(variable_map, current_node.id)
    # 其他类型（数字、表达式等）可按需扩展处理逻辑
    else:
        return None

# 示例代码
sample_code = """
pg_connection0 = "airflow_db1"
pg_connection = pg_connection0
"""

# 解析AST并构建符号表
tree = ast.parse(sample_code)
tracker = AssignmentTracker()
tracker.visit(tree)

# 获取目标变量的初始值并输出指定格式
initial_val = get_initial_value(tracker.variable_map, "pg_connection")
print(f'postgres_conn_id="{initial_val}"')

代码说明 #

• AssignmentTracker类遍历AST，将每个变量的赋值节点存入符号表variable_map；
• get_initial_value函数递归解析变量：若当前节点是字符串常量则直接返回值，若是变量名则继续追踪该变量的赋值；
• 运行示例代码后，会输出postgres_conn_id="airflow_db1"，完全符合需求。

扩展说明 #

如果需要处理更复杂场景，可针对性扩展逻辑：

• 多变量赋值：添加对ast.Tuple类型目标的处理；
• 表达式赋值：比如a = b + c，可根据需求决定是否追踪到最底层的常量；
• 条件赋值：需结合控制流分析，可借助astunparse等库辅助解析。

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