Python AST自动移除print/head/show/to_html无用代码

目录

• 一、使用 AST 的优势
• 二、环境准备
• 三、目标效果
• 四、完整可运行代码（⭐重点）
• 五、运行前后对比
• 清洗前
• 清洗后
• 六、核心原理解析（简要）
• 七、适用场景
• 八、可扩展方向（进阶）
• 九、结语

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

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

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

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

方式	问题
正则字符串	易误删、难维护
AST 抽象语法树	语法级安全、精准删除

pip install astor

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

删除以下代码：

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

复制即可运行

import ast
import astor


class MethodCallRemover(ast.NodeTransformer):
    """
    AST 修改器：
    1. 删除指定方法调用（如 head / show / to_html）
    2. 删除整行 print(...) 代码
    """

    def __init__(self, method_names):
        self.method_names = method_names

    def visit_Call(self, node):
        """
        删除嵌套的方法调用，例如：
        data.to_html()
        df.head()
        """
        if isinstance(node.func, ast.Attribute):
            if node.func.attr in self.method_names:
                return None
        return self.generic_visit(node)

    def remove_print_and_show(self, tree):
        """
        删除顶层的 print / show / head 语句（整行）
        """
        new_body = []

        for node in tree.body:
            # 删除 print(...)
            if (
                isinstance(node, ast.Expr)
                and isinstance(node.value, ast.Call)
                and isinstance(node.value.func, ast.Name)
                and node.value.func.id == "print"
            ):
                continue

            # 删除 obj.show() / obj.head()
            if (
                isinstance(node, ast.Expr)
                and isinstance(node.value, ast.Call)
                and isinstance(node.value.func, ast.Attribute)
                and node.value.func.attr in self.method_names
            ):
                continue

            new_body.append(node)

        tree.body = new_body
        return tree


# ===================== 测试代码（示例） =====================

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. 解析为 AST
tree = 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 中关键节点：

节点	含义
ast.Expr	独立的一行
ast.Call	函数/方法调用
ast.Attribute	obj.method

只要识别出：

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

就能 整行删除。

Notebook 转生产脚本

AI 生成代码自动去噪

pandas / matplotlib 批量清洗

自动化代码重构

企业代码规范化

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

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

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

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

到此这篇关于Python AST自动移除print/head/show/to_html无用代码的文章就介绍到这了,更多相关Python AST移除无用代码内容请搜索风君子博客以前的文章或继续浏览下面的相关文章希望大家以后多多支持风君子博客！

您可能感兴趣的文章:

• 基于Python AST实现代码安全检测功能
• 浅析AST抽象语法树及Python代码实现
• Python AST 模块实战演示
• python ast模块详析与用法
• Python Ast抽象语法树的介绍及应用详解