-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph.py
67 lines (54 loc) · 1.62 KB
/
graph.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
from typing import Set, Optional
from pathlib import Path
from secrets import token_urlsafe
from pyvis import network as net
from models import User, Relation
from services import get_exchange_bounds
from utils import normalize_relations
from data import users, relations
from config import HTML_DIR
graph_options = """
var options = {
"autoResize": true,
"edges": {
"color": {
"color": "#848484",
"highlight": "#FF0000"
}
},
"physics": {
"barnesHut": {
"springConstant": 0,
"damping": 0.5,
"avoidOverlap": 0.75
},
"maxVelocity": 10,
"minVelocity": 0.75
}
}
"""
def html_graph(
nodes: Set[User] = users.copy(),
edges: Set[Relation] = relations.copy(),
user: Optional[str] = None,
) -> Path:
if user:
node_edges = {r for r in edges if user in r}
print(f"Removed {len(edges) - len(node_edges)} edges")
edges = normalize_relations(node_edges, get_exchange_bounds(node_edges), (1, 50))
nodes = {u for u in nodes if any(u in r for r in edges)}
g = net.Network(width="100%", height="100%")
g.set_options(graph_options)
# g.show_buttons(filter_=['physics'])
for node in nodes:
g.add_node(node.id, node.username, size=node.size)
for edge in edges:
try:
g.add_edge(edge.from_user.id, edge.to_user.id, width=edge.exchange)
except AssertionError as e:
print(e)
print(f"Nodes in graph: {len(g.get_nodes())}")
print(f"Edges in graph: {len(g.get_edges())}")
filename = HTML_DIR / f"{token_urlsafe(16)}.html"
g.save_graph(str(filename))
return filename