PrivacyScanData/utils/splitfile.py

import ast
import os

from utils.fileio import walk_files


def write_content_to_file(content, file_path):
    if not os.path.exists(file_path[:file_path.rindex("/")]):
        os.makedirs(file_path[:file_path.rindex("/")])
    with open(file_path, 'w') as file_out:
        file_out.write(content)


def split_file(file_dir, output_dir, endpoint=".py"):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for file_path in walk_files(file_dir, endpoint):
        with open(file_path, 'r') as file:
            content = file.read()
        try:
            root = ast.parse(content)
        except Exception as e:
            print("错误明细：", e.__class__.__name__, e, file_path)
            continue
        file_id = 1
        for node in root.body:
            if isinstance(node, ast.FunctionDef):
                new_file_name = file_path.replace(file_dir, output_dir).replace(".py", "_" + str(file_id) + ".py")
                write_content_to_file(ast.get_source_segment(content, node), new_file_name)
                file_id += 1
            elif isinstance(node, ast.ClassDef):
                for son in node.body:
                    if isinstance(son, ast.FunctionDef):
                        new_file_name = file_path.replace(file_dir, output_dir).replace(".py",
                                                                                        "_" + str(file_id) + ".py")
                        write_content_to_file(ast.get_source_segment(content, son), new_file_name)
                        file_id += 1


if __name__ == '__main__':
    split_file("/Users/liufan/program/PYTHON/sap2nd/Data/archive/new2", "../target/archive/new2")