python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > Python AST移除无用代码

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

作者:weixin_46244623

这篇文章主要为大家详细介绍了Python如何AST自动移除print/head/show/to_html等无用代码,文中的示例代码讲解详细,感兴趣的小伙伴可以了解下

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

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

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

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

一、使用 AST 的优势

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

二、环境准备

pip install astor

三、目标效果

删除以下代码:

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.Attributeobj.method

只要识别出:

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

就能 整行删除

七、适用场景

Notebook 转生产脚本

AI 生成代码自动去噪

pandas / matplotlib 批量清洗

自动化代码重构

企业代码规范化

八、可扩展方向(进阶)

九、结语

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

当你开始用 AST,你会发现:

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

您可能感兴趣的文章:
阅读全文