nodetype.py 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. import ast
  2. import os
  3. current_path = os.path.dirname(__file__)
  4. def init_node_map():
  5. with open(current_path + "/astfile/_ast.py", 'r') as file:
  6. content = file.read()
  7. ast.parse(content)
  8. root = ast.parse(content)
  9. mp = dict()
  10. index = 1
  11. edge_mp = dict()
  12. edge_index = 1
  13. for node in root.body:
  14. if isinstance(node, ast.ClassDef):
  15. mp[node.name] = index
  16. index += 1
  17. for son in node.body:
  18. if isinstance(son, ast.Assign) and son.targets[0].id == '_fields':
  19. for field in son.value.dims:
  20. if not field.n in edge_mp:
  21. edge_mp[field.n] = edge_index
  22. edge_index += 1
  23. return mp, edge_mp
  24. class NodeType:
  25. node_map, edge_map = init_node_map()
  26. def write_to_txt():
  27. with open("/Users/liufan/program/PYTHON/sap2nd/GnnForPrivacyScan/data/node.txt", 'w') as file:
  28. for (node, t) in NodeType.node_map.items():
  29. file.write(node + " " + str(t) + "\n")
  30. with open("/Users/liufan/program/PYTHON/sap2nd/GnnForPrivacyScan/data/edge.txt", 'w') as file2:
  31. for (node, t) in NodeType.edge_map.items():
  32. file2.write(node + " " + str(t) + "\n")
  33. write_to_txt()