Graphify / binary_tree_generator.py
ZahirJS's picture
Update binary_tree_generator.py
8845dfd verified
import graphviz
import json
from tempfile import NamedTemporaryFile
import os
def generate_binary_tree_diagram(json_input: str, output_format: str) -> str:
"""
Generates a binary tree diagram from JSON input.
Args:
json_input (str): A JSON string describing the binary tree structure.
It must follow the Expected JSON Format Example below.
output_format (str): The output format for the generated diagram.
Supported formats: "png" or "svg"
Expected JSON Format Example:
{
"root": {
"id": "root",
"label": "80",
"left": {
"id": "node_35",
"label": "35",
"left": {
"id": "node_12",
"label": "12",
"left": {
"id": "node_8",
"label": "8",
"right": {
"id": "node_10",
"label": "10"
}
},
"right": {
"id": "node_28",
"label": "28",
"left": {
"id": "node_15",
"label": "15",
"right": {
"id": "node_22",
"label": "22",
"left": {
"id": "node_18",
"label": "18"
},
"right": {
"id": "node_25",
"label": "25"
}
}
},
"right": {
"id": "node_32",
"label": "32",
"left": {
"id": "node_30",
"label": "30"
}
}
}
},
"right": {
"id": "node_65",
"label": "65",
"right": {
"id": "node_72",
"label": "72",
"left": {
"id": "node_68",
"label": "68"
},
"right": {
"id": "node_76",
"label": "76",
"right": {
"id": "node_78",
"label": "78"
}
}
}
}
},
"right": {
"id": "node_120",
"label": "120",
"left": {
"id": "node_95",
"label": "95",
"left": {
"id": "node_88",
"label": "88",
"left": {
"id": "node_85",
"label": "85"
},
"right": {
"id": "node_92",
"label": "92"
}
},
"right": {
"id": "node_110",
"label": "110",
"left": {
"id": "node_105",
"label": "105",
"right": {
"id": "node_108",
"label": "108"
}
}
}
},
"right": {
"id": "node_180",
"label": "180",
"left": {
"id": "node_140",
"label": "140",
"right": {
"id": "node_165",
"label": "165",
"left": {
"id": "node_155",
"label": "155"
},
"right": {
"id": "node_170",
"label": "170",
"right": {
"id": "node_175",
"label": "175"
}
}
}
},
"right": {
"id": "node_220",
"label": "220",
"left": {
"id": "node_200",
"label": "200",
"left": {
"id": "node_190",
"label": "190"
},
"right": {
"id": "node_210",
"label": "210"
}
}
}
}
}
}
}
Returns:
str: The filepath to the generated image file.
"""
try:
if not json_input.strip():
return "Error: Empty input"
data = json.loads(json_input)
if 'root' not in data:
raise ValueError("Missing required field: root")
dot = graphviz.Digraph(
name='BinaryTree',
format='png',
graph_attr={
'rankdir': 'TB',
'splines': 'line',
'bgcolor': 'white',
'pad': '0.5',
'nodesep': '0.8',
'ranksep': '1.0'
}
)
base_color = '#BEBEBE'
def add_binary_tree_nodes(node, current_depth=0):
if not node:
return
node_id = node.get('id', f'node_{current_depth}')
node_label = node.get('label', 'Node')
lightening_factor = 0.06
if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
base_color_safe = '#BEBEBE'
else:
base_color_safe = base_color
base_r = int(base_color_safe[1:3], 16)
base_g = int(base_color_safe[3:5], 16)
base_b = int(base_color_safe[5:7], 16)
current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
current_r = min(255, current_r)
current_g = min(255, current_g)
current_b = min(255, current_b)
node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
font_color = 'black'
font_size = max(9, 14 - current_depth)
dot.node(
node_id,
node_label,
shape='circle',
style='filled,rounded',
fillcolor=node_color,
fontcolor=font_color,
fontsize=str(font_size),
width='0.8',
height='0.8'
)
left_child = node.get('left')
if left_child:
add_binary_tree_nodes(left_child, current_depth + 1)
left_id = left_child.get('id', f'node_{current_depth + 1}_left')
dot.edge(
node_id,
left_id,
color='#4a4a4a',
arrowsize='0.8'
)
right_child = node.get('right')
if right_child:
add_binary_tree_nodes(right_child, current_depth + 1)
right_id = right_child.get('id', f'node_{current_depth + 1}_right')
dot.edge(
node_id,
right_id,
color='#4a4a4a',
arrowsize='0.8'
)
add_binary_tree_nodes(data['root'], current_depth=0)
with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
dot.render(tmp.name, format=output_format, cleanup=True)
return f"{tmp.name}.{output_format}"
except json.JSONDecodeError:
return "Error: Invalid JSON format"
except Exception as e:
return f"Error: {str(e)}"