-
Notifications
You must be signed in to change notification settings - Fork 16
/
graph_dict.py
161 lines (136 loc) · 6.4 KB
/
graph_dict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import csv
import os
from graphipy.graph.graph_base import BaseGraph
# Graph object implemented by dict
class DictGraph(BaseGraph):
def __init__(self):
BaseGraph.__init__(self)
self._type = "dictionary"
# a dict of nodes and edges, key=label, value={id, node}
self.nodes = {}
self.edges = {}
self.path = os.getcwd() + "\\csv"
if not os.path.exists(self.path):
os.mkdir(self.path)
def graph_type(self):
return self._type
def export_all_csv(self, prefix):
""" Exports all dictionaries as CSV files and returns path to file """
# Create folders to export to
export_path = self.path + "\\" + prefix + "\\"
export_path_node = export_path + "nodes\\"
export_path_edge = export_path + "edges\\"
if not os.path.exists(export_path):
os.mkdir(export_path)
if not os.path.exists(export_path_node):
os.mkdir(export_path_node)
if not os.path.exists(export_path_edge):
os.mkdir(export_path_edge)
nodes = self.get_nodes()
edges = self.get_edges()
for key in nodes:
# create a csv file and write in data
with open(export_path_node + key + ".csv", "w", newline='', encoding="utf-8") as f:
w = csv.writer(f)
has_first_row = False
for node_key in nodes[key]:
# if title row has not been created
if has_first_row == False:
# create the title row
key_row = vars(nodes[key][node_key]).keys()
w.writerow([s for s in key_row])
has_first_row = True
# add value row to csv file
value_row = vars(nodes[key][node_key]).values()
w.writerow([s for s in value_row])
for key in edges:
# create a csv file and write in data
with open(export_path_edge + key + ".csv", "w", newline='', encoding="utf-8") as f:
w = csv.writer(f)
has_first_row = False
for edge_key in edges[key]:
# if title row has not been created
if has_first_row == False:
# create the title row
key_row = vars(edges[key][edge_key]).keys()
w.writerow([s for s in key_row])
has_first_row = True
# add value row to csv file
value_row = vars(edges[key][edge_key]).values()
w.writerow([s for s in value_row])
return export_path
# export data frame to csv specified by node label and edge label
def export_csv(self, prefix, node_option=set(), edge_option=set()):
""" Exports specific nodes and edges as CSV files and returns path to file """
# Create folders to export to
export_path = self.path + "\\" + prefix + "\\"
export_path_node = export_path + "nodes\\"
export_path_edge = export_path + "edges\\"
if not os.path.exists(export_path):
os.mkdir(export_path)
if not os.path.exists(export_path_node):
os.mkdir(export_path_node)
if not os.path.exists(export_path_edge):
os.mkdir(export_path_edge)
if len(node_option) > 0:
# get all nodes
nodes = self.get_nodes()
for key in nodes.keys():
# if matches node label that user wants
if key in node_option:
# create a csv file and write in data
with open(export_path_node + key + ".csv", "w", newline='', encoding="utf-8") as f:
w = csv.writer(f)
has_first_row = False
# iterate through {id, node}
for node_key in nodes[key]:
# if title row has not been created
if has_first_row == False:
# create the title row
key_row = vars(nodes[key][node_key]).keys()
w.writerow([s for s in key_row])
has_first_row = True
# add value row to csv file
value_row = vars(nodes[key][node_key]).values()
w.writerow([s for s in value_row])
if len(edge_option) > 0:
# get all edges
edges = self.get_edges()
for key in edges.keys():
# if matches edge label that user wants
if key in edge_option:
# create a csv file and write in data
with open(export_path_edge + key + ".csv", "w", newline='', encoding="utf-8") as f:
w = csv.writer(f)
has_first_row = False
# iterate through {id, node}
for edge_key in edges[key]:
# if title row has not been created
if has_first_row == False:
key_row = vars(edges[key][edge_key]).keys()
w.writerow([s for s in key_row])
has_first_row = True
# add value row to csv file
value_row = vars(edges[key][edge_key]).values()
w.writerow([s for s in value_row])
return export_path
def create_node(self, node):
node_label = node.get_label_attribute()
node_id = node.get_id()
# if no dict for this label in nodes{}, create a new dict with key=label
if node_label not in self.nodes:
self.nodes[node_label] = {}
# add node to {id, node} in {label, {id, node}}
self.nodes[node_label][node_id] = node
def create_edge(self, edge):
edge_label = edge.get_label_attribute()
edge_id = edge.get_id()
# if no dict for this label in edges{}, create a new dict with key=label
if edge_label not in self.edges:
self.edges[edge_label] = {}
# add edge to {id, edge} in {label, {id, edge}}
self.edges[edge_label][edge_id] = edge
def get_nodes(self):
return self.nodes
def get_edges(self):
return self.edges