CodeWalk

AST模块:解析、分析和转换Python源码

作者:屠龙少年 · 2026-05-30 12:55

请介绍Python ast模块的核心功能和使用方法。如何将源码解析为AST(ast.parse)、遍历AST节点(ast.NodeVisitorast.NodeTransformer)以及将AST还原为代码(ast.unparse / compile + exec)。给出一个实际应用场景:自动添加函数日志或检测未使用的import。

回答

屠龙少年

AST基础操作

import ast

code = """
def add(a, b):
    return a + b
"""

tree = ast.parse(code)  # 解析为AST
print(ast.dump(tree, indent=2))  # 查看AST结构

# 编译AST为code object并执行
compiled = compile(tree, filename='<ast>', mode='exec')
exec(compiled)

NodeVisitor遍历

class FuncCallVisitor(ast.NodeVisitor):
    def __init__(self):
        self.calls = []
    
    def visit_Call(self, node):
        if isinstance(node.func, ast.Name):
            self.calls.append(node.func.id)
        self.generic_visit(node)

visitor = FuncCallVisitor()
visitor.visit(ast.parse(code))
print(visitor.calls)  # 输出所有函数调用

NodeTransformer自动添加日志

class AddLoggingTransformer(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        log_call = ast.Expr(
            value=ast.Call(
                func=ast.Name(id='print', ctx=ast.Load()),
                args=[ast.Constant(value=f"Entering {node.name}")],
                keywords=[]
            )
        )
        node.body.insert(0, log_call)
        return self.generic_visit(node)

tree = ast.parse(code)
transformer = AddLoggingTransformer()
new_tree = transformer.visit(tree)
ast.fix_missing_locations(new_tree)  # 修复位置信息
new_code = ast.unparse(new_tree)
print(new_code)

检测未使用import: 遍历ast.Import/ast.ImportFrom收集导入名,再遍历ast.Name检查引用,未被引用的import可报告删除。

工具生态astor(更完善的unparse)、libCST(Facebook的CST工具,保留注释和格式)、typed_ast(Python 2兼容)。