GnnForPrivacyScan/graph/astgraph/AstGraph.py

import ast
import os
import random

from graph.common.nodetype import *
from graph.common.graphtype import *
from utils.fileio import *

CUR_PATH = os.path.dirname(__file__)
i = 0


def ast_visit(node, graph=None, id_list=None):
    if graph is None:
        graph = list()
    if id_list is None:
        id_list = [node]
    iter_field = ast.iter_fields(node)
    for field, value in iter_field:
        # print(type(node).__name__, field, value)
        if isinstance(value, list):
            for item in value:
                if isinstance(item, ast.AST):
                    id_list.append(item)
                    graph.append(
                        (id_list.index(node) + 1, NodeType.node_map[type(node).__name__], 1,
                         id_list.index(item) + 1, NodeType.node_map[type(item).__name__])
                        )
                    ast_visit(item, graph, id_list)
        elif isinstance(value, ast.AST):
            id_list.append(value)
            graph.append(
                (id_list.index(node) + 1, NodeType.node_map[type(node).__name__], 1,
                 id_list.index(value) + 1, NodeType.node_map[type(value).__name__])
            )
            ast_visit(value, graph, id_list)


def gen_graph_from_file(file_path):
    with open(file_path, 'r') as file:
        content = file.read()
    func_graph = list()
    try:
        root = ast.parse(content)
        for node in root.body:
            if isinstance(node, ast.FunctionDef):
                ast_visit(node, func_graph)

    except IndentationError:
        print("IndentationError: ", file_path)
    except SyntaxError:
        print("SyntaxError: ", file_path)
    except:
        print("other: ", file_path)

    return func_graph


def gen_graph_to_txt():
    path_ori = CUR_PATH + "/../../data/purposeSplit"
    kinds = os.listdir(path_ori)
    kinds.remove('.DS_Store')
    for kind in kinds:
        graph_type = GraphType.type[kind]
        path_out_train = CUR_PATH + "/../../data/traindata/train/" + str(kind) + ".txt"
        path_out_test = CUR_PATH + "/../../data/traindata/test/" + str(kind) + ".txt"
        file_path_list = walk_files(path_ori + "/" + kind)
        random.shuffle(file_path_list)
        lens = int(len(file_path_list) / 4 * 3)
        with open(path_out_train, 'w') as file:
            for file_path in file_path_list[:lens]:
                func_graph = gen_graph_from_file(file_path)
                for edge in func_graph:
                    file.write(
                        "" + str(edge[0]) + " " + str(edge[1]) + " " + str(edge[2]) + " " + str(edge[3]) + " " + str(
                            edge[
                                4]) + "\n")
                file.write("? " + str(graph_type) + " " + file_path.replace(path_ori + "/", "") + "\n\n")
        with open(path_out_test, 'w') as file:
            for file_path in file_path_list[lens: len(file_path_list)]:
                func_graph = gen_graph_from_file(file_path)
                for edge in func_graph:
                    file.write(
                        "" + str(edge[0]) + " " + str(edge[1]) + " " + str(edge[2]) + " " + str(edge[3]) + " " + str(
                            edge[
                                4]) + "\n")
                file.write("? " + str(graph_type) + " " + file_path.replace(path_ori + "/", "") + "\n\n")


if __name__ == '__main__':
    gen_graph_to_txt()
    # graph = gen_graph_from_file("/Users/liufan/program/PYTHON/sap2nd/GnnForPrivacyScan/data/purposeSplit/Directory/advance_touch_1.py", )
    # a = 5