ZahirJS commited on
Commit
31c7f3c
·
verified ·
1 Parent(s): 6eddc8a

Create graph_generator_utils.py

Browse files
Files changed (1) hide show
  1. graph_generator_utils.py +79 -0
graph_generator_utils.py ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import graphviz
2
+
3
+ def add_nodes_and_edges(dot: graphviz.Digraph, parent_id: str, nodes_list: list, current_depth: int, base_color: str):
4
+ """
5
+ Recursively adds nodes and edges to a Graphviz Digraph object,
6
+ applying a color gradient and consistent styling.
7
+
8
+ Args:
9
+ dot (graphviz.Digraph): The Graphviz Digraph object to modify.
10
+ parent_id (str): The ID of the parent node for the current set of nodes.
11
+ nodes_list (list): A list of dictionaries, each representing a node
12
+ with 'id', 'label', 'relationship', and optional 'subnodes'.
13
+ current_depth (int): The current depth in the graph hierarchy (0 for central node).
14
+ base_color (str): The hexadecimal base color for the deepest nodes.
15
+ """
16
+ # Calculate color for current depth, making it lighter
17
+ # This factor determines how quickly the color lightens per level.
18
+ lightening_factor = 0.12
19
+
20
+ # Convert base_color hex to RGB for interpolation
21
+ base_r = int(base_color[1:3], 16)
22
+ base_g = int(base_color[3:5], 16)
23
+ base_b = int(base_color[5:7], 16)
24
+
25
+ # Calculate current node color by blending towards white
26
+ current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
27
+ current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
28
+ current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
29
+
30
+ # Clamp values to 255 to stay within valid RGB range
31
+ current_r = min(255, current_r)
32
+ current_g = min(255, current_g)
33
+ current_b = min(255, current_b)
34
+
35
+ node_fill_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
36
+
37
+ # Font color: white for dark nodes, black for very light nodes for readability
38
+ font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
39
+
40
+ # Edge colors and font sizes
41
+ edge_color = '#4a4a4a' # Dark gray for lines
42
+ # Font size adjusts based on depth, ensuring a minimum size
43
+ font_size = max(9, 14 - (current_depth * 2))
44
+ edge_font_size = max(7, 10 - (current_depth * 1))
45
+
46
+ for node in nodes_list:
47
+ node_id = node.get('id')
48
+ label = node.get('label')
49
+ relationship = node.get('relationship')
50
+
51
+ # Basic validation for node data
52
+ if not all([node_id, label, relationship]):
53
+ raise ValueError(f"Invalid node: {node}")
54
+
55
+ # Add node with specified style
56
+ dot.node(
57
+ node_id,
58
+ label,
59
+ shape='box', # All nodes are rectangular
60
+ style='filled,rounded', # Filled and rounded corners
61
+ fillcolor=node_fill_color,
62
+ fontcolor=font_color,
63
+ fontsize=str(font_size)
64
+ )
65
+
66
+ # Add edge from parent to current node
67
+ dot.edge(
68
+ parent_id,
69
+ node_id,
70
+ label=relationship,
71
+ color=edge_color,
72
+ fontcolor=edge_color, # Edge label color also dark gray
73
+ fontsize=str(edge_font_size)
74
+ )
75
+
76
+ # Recursively call for subnodes
77
+ if 'subnodes' in node:
78
+ add_nodes_and_edges(dot, node_id, node['subnodes'], current_depth + 1, base_color)
79
+