Create graph_generator_utils.py
Browse files- 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 |
+
|