Python AST 实战：自动移除 print / head / show / to_html 等无用代码行

在数据分析、Notebook 转生产代码、AI 生成代码清洗等场景中，我们经常需要：

自动删除print()、DataFrame.head()、plt.show()、to_html()等仅用于展示的代码，而不影响业务逻辑

正则不可靠，AST 才是王道。

本文将通过一个完整可运行示例，教你如何使用Python AST 对源码进行结构级修改。

一、使用 AST 的优势

二、环境准备

pipinstallastor

• ast：Python 内置模块
• astor：将 AST 转回源码

三、目标效果

删除以下代码：

print(...)df.head()plt.show()data.to_html()

四、完整可运行代码（⭐重点）

✅复制即可运行

importastimportastorclassMethodCallRemover(ast.NodeTransformer):""" AST 修改器： 1. 删除指定方法调用（如 head / show / to_html） 2. 删除整行 print(...) 代码 """def__init__(self,method_names):self.method_names=method_namesdefvisit_Call(self,node):""" 删除嵌套的方法调用，例如： data.to_html() df.head() """ifisinstance(node.func,ast.Attribute):ifnode.func.attrinself.method_names:returnNonereturnself.generic_visit(node)defremove_print_and_show(self,tree):""" 删除顶层的 print / show / head 语句（整行） """new_body=[]fornodeintree.body:# 删除 print(...)if(isinstance(node,ast.Expr)andisinstance(node.value,ast.Call)andisinstance(node.value.func,ast.Name)andnode.value.func.id=="print"):continue# 删除 obj.show() / obj.head()if(isinstance(node,ast.Expr)andisinstance(node.value,ast.Call)andisinstance(node.value.func,ast.Attribute)andnode.value.func.attrinself.method_names):continuenew_body.append(node)tree.body=new_bodyreturntree# ===================== 测试代码（示例） =====================code=""" import pandas as pd import matplotlib.pyplot as plt data = pd.DataFrame({ 'A': [1, 2, 3], 'B': [4, 5, 6] }) print("原始数据：") print(data) data.head() plt.plot(data['A'], data['B']) plt.title("测试图表") plt.show() html = data.to_html() print(html) """# ===================== AST 处理流程 =====================# 1. 解析为 ASTtree=ast.parse(code)# 2. 要删除的方法名methods_to_remove=["head","show","render","render_notebook","to_html",]# 3. 执行清洗remover=MethodCallRemover(methods_to_remove)new_tree=remover.remove_print_and_show(tree)# 4. 转回源码clean_code=astor.to_source(new_tree)# 5. 输出结果print("====== 清洗后的代码 ======")print(clean_code)

五、运行前后对比

❌ 清洗前

print(data)data.head()plt.show()data.to_html()

✅ 清洗后

data=pd.DataFrame(...)plt.plot(...)plt.title("测试图表")

🎯只保留业务逻辑，彻底移除展示代码

六、核心原理解析（简要）

AST 中关键节点：

只要识别出：

Expr(value=Call(func=Attribute(attr="head")))

就能整行删除。

七、适用场景

✅ Notebook 转生产脚本
✅ AI 生成代码自动去噪
✅ pandas / matplotlib 批量清洗
✅ 自动化代码重构
✅ 企业代码规范化

八、可扩展方向（进阶）

• 🔹 自动删除所有plt.*
• 🔹 AST 自动加日志
• 🔹 AST + LLM 代码修复
• 🔹 批量处理.py文件

九、结语

AST 是 Python 工程师的“代码手术刀”

当你开始用 AST，你会发现：

• 正则 = 土办法
• AST = 工程级解决方案