python删除一个函数 ast

python删除一个函数 ast

import ast
import sys


class FunctionRemover(ast.NodeTransformer):
    def __init__(self, func_name):
        self.func_name = func_name

    def visit_FunctionDef(self, node):
        # 如果是我们要删除的函数，添加占位标志并返回None来从树中移除它
        if node.name == self.func_name:
            # print("node",node)
            # print("node.name ",node.name )
            # 创建一个占位标志的注释节点
            placeholder = ast.Expr(value=ast.Constant(value="# Function removed: " + node.name))
            # return [placeholder, None]
            # return [placeholder, placeholder]
            return placeholder
        return node

    def visit_Call(self, node):
        # 检查函数调用是否是我们想要删除的函数
        if isinstance(node.func, ast.Name) and node.func.id == self.func_name:
            # 创建一个占位标志的注释节点
            placeholder = ast.Expr(value=ast.Constant(value="# Function call removed: " + self.func_name))
            return placeholder
        return node



# def remove_function_from_code(code, func_name):
#     tree = ast.parse(code)
#     # AttributeError: 'Assign' object has no attribute 'value'
#     remover = FunctionRemover(func_name)
#     new_tree = remover.visit(tree)
#     new_code = ast.unparse(new_tree)
#     return new_code

def remove_function_from_code(code, func_name):
    tree = ast.parse(code)
    remover = FunctionRemover(func_name)
    new_tree = remover.visit(tree)
    ast.fix_missing_locations(new_tree)  # 修复AST树的位置信息
    new_code = ast.unparse(new_tree)
    return new_code



import ast

class FunctionNameFinder(ast.NodeVisitor):
    def __init__(self):
        self.function_names = []

    def visit_FunctionDef(self, node):
        # 将函数名添加到列表中
        self.function_names.append(node.name)
        # 继续遍历AST
        self.generic_visit(node)

# def list_function_names(file_path):
#     with open(file_path, "r") as file:
#         code = file.read()
#     tree = ast.parse(code)
#     visitor = FunctionNameFinder()
#     visitor.visit(tree)
#     return visitor.function_names

def list_function_names(code):
    # with open(file_path, "r") as file:
    #     code = file.read()
    tree = ast.parse(code)
    visitor = FunctionNameFinder()
    visitor.visit(tree)
    return visitor.function_names


from top.starp.util import file_util

import random


# 使用示例
# file_path = "example.py"  # 将这里的路径替换为你的Python文件路径
# function_names = list_function_names(file_path)
# print("函数名列表:", function_names)
# file_path=rf"D:\proj\python\data_enhancement\parquet_language_percent.py"



def code_random_rm_func(code):
    function_names = list_function_names(code)
    # 使用random.choice随机选取一个元素
    func_name = random.choice(function_names)

    new_code = remove_function_from_code(code, func_name)
    # print(new_code)

    # print(new_code)
    # 补充这里的代码  # Function removed:

    return new_code



def code_path_random_rm_func(file_path=rf"D:\proj\python\data_enhancement\gen_sentence.py"):
    # file_path=rf"D:\proj\python\data_enhancement\parquet_code_desp.py"

    code=file_util.read_file(file_path)
    return code_random_rm_func(code)