如何用Python AST遍历代码提取函数名与顺序化返回语句？

问题描述 #

我有一个Python脚本，想要对其进行“解码”以便转换为XML，但目前正尝试理解如何使用ast.walk()以及如何从中提取所需信息。我在simple_if.py中编写了如下简单函数，目标是提取函数名和按代码逻辑顺序排列的返回语句列表：

def if_else(abc):
    if abc > 0:
        return "Hello"
    elif abc < 0:
        return "Goodbye"
    return "Neither Hello Nor Goodbye"

我将上述代码放入文件后，运行以下AST相关代码：

with open("simple_if.py", "r", encoding = "utf-8") as ast_tree_walker
    tree = ast.parse(ast_tree_walker.read())

expList = []
counter = 0

for node in ast.walk(tree):
    print(node)
    if isInstance(node, ast.FunctionDef) and counter == 0:
        expList.append(node.__dict__["name"]
        counter = 1 # this adds the functiona name if_else to the list 
    print(node.__dict__)

上述代码会在终端输出节点、字典和参数，但目前仅能提取函数名（因其是树中第一个非Module节点）。我知道所需信息已输出，但不清楚如何保持返回语句的顺序——例如"Neither Hello Nor Goodbye"会出现在"Hello"或"Goodbye"之前，推测这是由于树结构中最后的return与if语句处于同一层级，但不确定是否能维持代码逻辑顺序。我想知道如何得到列表["if_else", "Hello", "Goodbye", "Neither Hello Nor Goodbye"]？

此外还有细节问题：上述代码会输出节点类型，如<ast.Module object at ...>、<ast.FunctionDef object at ...>、<ast.If object at ...>、<ast.Return object at ...>... 当遇到<ast.If object at ...>节点时，我应该直接访问其'body'键，还是等待子节点被访问后再提取信息？

ast.If节点的子节点信息如下：

#children of ast.If
{'test': <ast.Compare object at>, 'body': <ast.Return object at >, 'orelse': [<ast.If object at>], 'lineno':2, 'col_offset':7}

解决方案 #

1. 为什么ast.walk()会打乱顺序 #

ast.walk()采用深度优先遍历，遍历顺序是树的结构顺序，并非代码的逻辑执行顺序。比如你的代码里，最后一个return和If节点同属函数体的一级节点，walk的遍历顺序可能先捕获到这个return，再遍历If节点内部的返回语句，导致顺序混乱。

2. 按逻辑顺序提取的实现代码 #

要严格遵循代码逻辑顺序提取返回语句，需要手动递归遍历函数体，按分支顺序处理节点：

import ast

def extract_target_info(tree):
    result = []
    # 定位到目标函数
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            result.append(node.name)
            # 递归遍历函数体，按逻辑顺序处理节点
            def traverse_statements(stmts):
                for stmt in stmts:
                    if isinstance(stmt, ast.Return):
                        # 提取字符串类型的返回值
                        if isinstance(stmt.value, ast.Constant) and isinstance(stmt.value.value, str):
                            result.append(stmt.value.value)
                    elif isinstance(stmt, ast.If):
                        # 先处理if分支代码
                        traverse_statements(stmt.body)
                        # 再处理else/elif分支代码
                        traverse_statements(stmt.orelse)
            traverse_statements(node.body)
            # 仅处理第一个函数，如需多函数可删除break
            break
    return result

# 读取文件并解析AST
with open("simple_if.py", "r", encoding="utf-8") as f:
    ast_tree = ast.parse(f.read())

# 获取结果
final_list = extract_target_info(ast_tree)
print(final_list)  # 输出: ['if_else', 'Hello', 'Goodbye', 'Neither Hello Nor Goodbye']

3. If节点的处理方式 #

不要等待ast.walk()自动遍历子节点，必须主动访问If节点的body和orelse属性，先处理body（if满足时执行的分支），再处理orelse（elif/else分支）。这样才能保证返回语句的顺序和代码书写、执行的逻辑完全一致。依赖ast.walk()的自动遍历无法控制顺序，必然会出现混乱。

补充说明 #

• 代码默认只处理第一个函数，若需要处理多个函数，可删除break，并将结果改为字典结构（如{函数名: 返回值列表}）
• 当前仅支持提取字符串常量类型的返回值，若返回值是变量、表达式等，可根据需求扩展traverse_statements函数的处理逻辑

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