AST模块:解析、分析和转换Python源码
请介绍Python ast模块的核心功能和使用方法。如何将源码解析为AST(ast.parse)、遍历AST节点(ast.NodeVisitor、ast.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兼容)。