-
Notifications
You must be signed in to change notification settings - Fork 6.7k
/
visualizer.py
82 lines (74 loc) · 1.9 KB
/
visualizer.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
from pyvis.network import Network
import os
class Visualizer:
def __init__(self):
self.directory = "generated"
if not os.path.exists(self.directory):
os.makedirs(self.directory)
def render(self, query_lineage_response, scenario_name):
net = self.get_network()
for vertex in query_lineage_response["Vertices"]:
arn = vertex["Arn"]
if "Type" in vertex:
label = vertex["Type"]
else:
label = None
lineage_type = vertex["LineageType"]
name = self.get_name(arn)
title = self.get_title(arn)
net.add_node(vertex["Arn"], label=label+"\n"+lineage_type, title=title, shape="circle")
for edge in query_lineage_response["Edges"]:
source = edge["SourceArn"]
dest = edge["DestinationArn"]
net.add_edge(dest, source)
return net.show(f"{self.directory}/{scenario_name}.html")
def get_title(self, arn):
return f"Arn: {arn}"
def get_name(self, arn):
name = arn.split("/")[1]
return name
def get_network(self):
net = Network(height="400px", width="800px", directed=True, notebook=True)
net.set_options(
"""
var options = {
"nodes": {
"borderWidth": 3,
"shadow": {
"enabled": true
},
"shapeProperties": {
"borderRadius": 3
},
"size": 11,
"shape": "circle"
},
"edges": {
"arrows": {
"to": {
"enabled": true
}
},
"color": {
"inherit": true
},
"smooth": false
},
"layout": {
"hierarchical": {
"enabled": true,
"direction": "LR",
"sortMethod": "directed"
}
},
"physics": {
"hierarchicalRepulsion": {
"centralGravity": 0
},
"minVelocity": 0.75,
"solver": "hierarchicalRepulsion"
}
}
"""
)
return net