Lauri Koskenniemi
cff261519f
Implemented working version of Graph. Issues to solve remain, such as ordering of nodes in a block.
93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
import graphviz
|
|
import uuid
|
|
|
|
def get_id():
|
|
return "id" + str(uuid.uuid4().hex)
|
|
|
|
# One cluster of nodes per block
|
|
class Graph():
|
|
def __init__(self):
|
|
self.dot = None
|
|
self.links = {}
|
|
|
|
def create_graph(self, config, data):
|
|
self.dot = graphviz.Graph('testgraph', graph_attr={'center': 'true', 'compound': 'true', 'nodesep': '0.2', 'ranksep': '0.5 equally'}, node_attr={'shape': 'plaintext'})
|
|
self.dot.format = 'svg'
|
|
|
|
blocks = data.get('blocks', [])
|
|
for block in blocks:
|
|
texts = block.get('texts', [])
|
|
center_node = int(len(texts) / 2)
|
|
cluster_id = 'cluster_' + get_id()
|
|
|
|
with self.dot.subgraph(name=cluster_id) as cluster:
|
|
prev_node_id = ""
|
|
for text in texts:
|
|
node_id = get_id()
|
|
cluster.node(node_id, label=text.get('text', ''))
|
|
if prev_node_id:
|
|
cluster.edge(prev_node_id, node_id, style='invis', rank='same')
|
|
prev_node_id = node_id
|
|
for link in text.get('links', []):
|
|
self.links[link] = node_id
|
|
|
|
if center_node == 0:
|
|
for link in block.get('links', []):
|
|
if link in self.links:
|
|
self.dot.edge(self.links[link], node_id, lhead=cluster_id)
|
|
del self.links[link]
|
|
center_node -= 1
|
|
|
|
def export_graph(self, out_dir):
|
|
self.dot.render(directory=out_dir)
|
|
|
|
def __str__(self):
|
|
return self.dot.source
|
|
|
|
|
|
# One node per block
|
|
class RecordGraph:
|
|
def __init__(self):
|
|
self.dot = None
|
|
self.edges = []
|
|
self.node_counter = 0
|
|
|
|
def create_graph(self, config, data):
|
|
# TODO: process config, use for graph init
|
|
self.dot = graphviz.Graph('testgraph', graph_attr={'center': 'true', 'compound': 'true'}, node_attr={'shape': 'record'})
|
|
self.dot.format = 'svg' # TODO: move to export
|
|
|
|
blocks = data.get('blocks', [])
|
|
for block in blocks:
|
|
self.node_counter += 1
|
|
|
|
# Create node id
|
|
node_id = f'nr{self.node_counter}'
|
|
links = block.get('links', [])
|
|
if len(links) > 0:
|
|
node_id = f'n{links[0]}'
|
|
|
|
# Create text table for node
|
|
texts = []
|
|
for text in block.get('texts', []):
|
|
person = text.get('text', '').replace('\n', '\\n')
|
|
link = text.get('links', [])
|
|
if len(link) > 0:
|
|
person = fr'<l{link[0]}> {person}'
|
|
self.edges.append((fr'{node_id}:l{link[0]}', fr'n{link[0]}'))
|
|
texts.append(person)
|
|
|
|
# Add node
|
|
table = r"|".join(texts)
|
|
self.dot.node(node_id, table)
|
|
|
|
# Add edges
|
|
self.dot.edges(self.edges)
|
|
|
|
def export_graph(self, out_dir):
|
|
self.dot.render(directory=out_dir)
|
|
|
|
|
|
def __str__(self):
|
|
return self.dot.source
|