Create binary_tree_generator.py

binary_tree_generator.py

@@ -0,0 +1,141 @@
+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.
+
+    Expected JSON Format Example:
+    {
+      "root": {
+        "id": "root",
+        "label": "50",
+        "left": {
+          "id": "left_1",
+          "label": "30",
+          "left": {
+            "id": "left_2",
+            "label": "20"
+          },
+          "right": {
+            "id": "right_2",
+            "label": "40"
+          }
+        },
+        "right": {
+          "id": "right_1",
+          "label": "70",
+          "left": {
+            "id": "left_3",
+            "label": "60"
+          },
+          "right": {
+            "id": "right_3",
+            "label": "80"
+          }
+        }
+      }
+    }
+
+    Returns:
+        str: The filepath to the generated PNG 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',        # Top-to-Bottom layout (vertical hierarchy)
+                'splines': 'line',      # Straight lines
+                'bgcolor': 'white',     # White background
+                'pad': '0.5',           # Padding around the graph
+                'nodesep': '0.8',       # Spacing between nodes
+                'ranksep': '1.0'        # Spacing between levels
+            }
+        )
+        
+        base_color = '#19191a' # Hardcoded base color
+        
+        def add_binary_tree_nodes(node, current_depth=0):
+            """
+            Add binary tree nodes recursively with proper styling.
+            """
+            if not node:
+                return
+            
+            node_id = node.get('id', f'node_{current_depth}')
+            node_label = node.get('label', 'Node')
+            
+            # Calculate color opacity based on depth
+            max_depth = 5  # Assume maximum depth for color calculation
+            if current_depth >= max_depth:
+                opacity = '80'  # Minimum opacity
+            else:
+                opacity_value = int(255 * (1.0 - (current_depth * 0.6 / max_depth)))
+                opacity = format(opacity_value, '02x')
+            
+            node_color = f"{base_color}{opacity}"
+            font_color = 'white' if current_depth < 3 else 'black'
+            
+            # Add the current node
+            dot.node(
+                node_id,
+                node_label,
+                shape='circle',
+                style='filled',
+                fillcolor=node_color,
+                fontcolor=font_color,
+                fontsize='14',
+                width='0.8',
+                height='0.8'
+            )
+            
+            # Process left child
+            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='#666666',
+                    arrowsize='0.8'
+                )
+            
+            # Process right child
+            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='#666666',
+                    arrowsize='0.8'
+                )
+        
+        # Add binary tree nodes and edges recursively
+        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)}"