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

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("测试图表")

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