def generate_graph(self):
"""
To be used after the data has been saved to an enrondump file by the parser
:return:
"""
graph = Graph('Enron')
emails = json.loads(open('enrondump').read())
for email in emails:
sender_email = email['from']
if graph.has_node_by_label(sender_email):
sender = graph.get_node_from_label(label=sender_email)
else:
sender = Node(label=sender_email)
graph.add_node(sender)
recipients = json.loads(email['To'])
for recipient_email in recipients:
if graph.has_node_by_label(recipient_email):
recipient = graph.get_node_from_label(recipient_email)
else:
recipient = Node(label=recipient_email)
graph.add_node(recipient)
if not graph.has_edge(sender, recipient):
graph.add_edge(sender, recipient, 1)
else:
# Increment the weight by 1
graph.update_edge(sender, recipient,
graph.get_weight(sender, recipient,) + 1)
return graph
def initialise_sample_graph(sample_graph):
graph = Graph(name='SampleGraph')
for each_node in sample_graph['nodes'].keys():
graph.add_node(Node(label=each_node))
for each_edge in sample_graph['weights'].keys():
node = graph.get_node_from_label(each_edge[0])
neighbor = graph.get_node_from_label(each_edge[1])
graph.add_edge(node, neighbor, sample_graph['weights'][each_edge])
return graph