Update app.py

app.py

@@ -1,19 +1,44 @@
-import os
-import logging
 import gradio as gr
 import pandas as pd
-import base64
+import networkx as nx
+import matplotlib.pyplot as plt
 import pdfkit
+import base64
 import shutil
-from graphviz import Digraph
-
-# Debug prints for environment variables
-print("PATH =", os.environ.get("PATH"))
-print("GRAPHVIZ_DOT =", os.environ.get("GRAPHVIZ_DOT"))
+import os
+import logging
 
-# Set up logging for debugging
 logging.basicConfig(level=logging.INFO)
 
+# -------------------------------
+# Helper: Custom tree layout
+# -------------------------------
+def hierarchy_pos(G, root, width=1.0, vert_gap=0.2, vert_loc=0, xcenter=0.5):
+    """
+    If the graph is a tree, this will return positions to plot a hierarchical layout.
+    
+    G: the tree (a DiGraph)
+    root: the root node of the tree
+    width: horizontal space allocated for the tree
+    vert_gap: gap between levels of the tree
+    vert_loc: vertical location of root
+    xcenter: horizontal location of root
+    """
+    def _hierarchy_pos(G, root, left, right, vert_loc, pos, parent=None):
+        pos[root] = ((left + right) / 2, vert_loc)
+        children = list(G.successors(root))
+        if children:
+            dx = (right - left) / len(children)
+            nextx = left
+            for child in children:
+                pos = _hierarchy_pos(G, child, nextx, nextx + dx, vert_loc - vert_gap, pos, root)
+                nextx += dx
+        return pos
+    return _hierarchy_pos(G, root, 0, width, vert_loc, {})
+
+# -------------------------------
+# Helpers for system dependencies
+# -------------------------------
 def get_wkhtmltopdf_path():
     path = shutil.which("wkhtmltopdf")
     if path:
@@ -23,6 +48,9 @@ def get_wkhtmltopdf_path():
             return candidate
     return None
 
+# -------------------------------
+# File processing helper
+# -------------------------------
 def process_uploaded_file(file):
     logging.info(f"Processing uploaded file: {file}")
     if isinstance(file, str):
@@ -33,7 +61,6 @@ def process_uploaded_file(file):
             logging.error(f"Provided file path '{file}' is a directory.")
             return None
         return file
-
     if isinstance(file, dict):
         if "data" in file and file["data"]:
             temp_path = "/tmp/uploaded.csv"
@@ -57,19 +84,19 @@ def process_uploaded_file(file):
             return None
     return None
 
+# -------------------------------
+# Main function: Generate Chart
+# -------------------------------
 def generate_chart(file, title):
     logging.info(f"generate_chart called with file: {file} and title: {title}")
     if not file:
         return None, "Please upload a CSV file."
-    
     file_path = process_uploaded_file(file)
     if not file_path:
         return None, "Error processing uploaded file. Please upload a valid CSV file."
-    
     if os.path.isdir(file_path):
         logging.error(f"Processed file path '{file_path}' is a directory.")
         return None, "Uploaded file is a directory. Please upload a valid CSV file."
-    
     try:
         df = pd.read_csv(file_path)
     except Exception as e:
@@ -80,35 +107,69 @@ def generate_chart(file, title):
     if not expected_columns.issubset(set(df.columns)):
         return None, "CSV must contain Name, Role, and Reporting To columns."
     
-    # Create a Graphviz Digraph for the org chart.
-    dot = Digraph(comment=title, format='png')
-    dot.attr(rankdir='TB')
-    dot.attr('node', shape='box', style='rounded,filled', fillcolor='lightblue',
-             fontsize='10', fontname='Helvetica', margin='0.1')
-    
-    # Add nodes.
+    # Build tree using names as nodes; store label as attribute
+    T = nx.DiGraph()
     for idx, row in df.iterrows():
-        node_id = row['Name']  # Assumes names are unique.
-        label = f"{row['Name']}\n({row['Role']})"
-        dot.node(node_id, label=label)
+        name = row["Name"]
+        role = row["Role"]
+        label = f"{name}\n({role})"
+        T.add_node(name, label=label)
     
-    # Add edges.
+    # Add edges (manager -> subordinate)
     for idx, row in df.iterrows():
-        if pd.notna(row['Reporting To']) and row['Reporting To'] != "":
-            manager = row['Reporting To']
-            subordinate = row['Name']
-            if manager not in df['Name'].values:
-                dot.node(manager, label=manager)
-            dot.edge(manager, subordinate)
-    
-    # Render diagram to PNG.
-    output_path = "/tmp/chart"
-    dot.render(output_path, cleanup=True)
-    chart_path = output_path + ".png"
-    
-    # Build HTML content with embedded chart image.
-    with open(chart_path, "rb") as img_file:
-        encoded = base64.b64encode(img_file.read()).decode()
+        if pd.notna(row["Reporting To"]) and row["Reporting To"] != "":
+            manager = row["Reporting To"]
+            subordinate = row["Name"]
+            if manager not in T.nodes:
+                # If manager not present, add with label = manager
+                T.add_node(manager, label=manager)
+            T.add_edge(manager, subordinate)
+    
+    # Determine the root (node with in-degree 0)
+    roots = [n for n, d in T.in_degree() if d == 0]
+    if not roots:
+        return None, "Could not determine root of the organization tree."
+    root = roots[0]
+    
+    pos = hierarchy_pos(T, root, width=1.0, vert_gap=0.2, vert_loc=1.0, xcenter=0.5)
+    
+    # Draw the tree using Matplotlib
+    plt.figure(figsize=(10, 8))
+    ax = plt.gca()
+    
+    # Draw nodes
+    for node, (x, y) in pos.items():
+        label = T.nodes[node].get("label", node)
+        lines = label.split("\n")
+        max_line_length = max(len(line) for line in lines)
+        # Use dynamic sizing: adjust factors as needed.
+        node_width = max_line_length * 3  # factor for width (in points)
+        node_height = len(lines) * 8      # factor for height (in points)
+        # For plotting, we use a conversion factor to scale node size relative to the plot.
+        # Here, we assume pos is normalized [0,1] and we scale sizes accordingly.
+        # Adjust these scaling factors as needed.
+        scale = 2  # adjust scale factor if needed
+        rect = plt.Rectangle((x - node_width/200.0*scale, y - node_height/200.0*scale),
+                             node_width/100.0*scale, node_height/100.0*scale,
+                             facecolor='lightblue', edgecolor='black', lw=1.5, zorder=2)
+        ax.add_patch(rect)
+        plt.text(x, y, label, horizontalalignment='center', verticalalignment='center', fontsize=8, zorder=3)
+    
+    # Draw edges as straight lines connecting centers
+    for u, v in T.edges():
+        if u in pos and v in pos:
+            x1, y1 = pos[u]
+            x2, y2 = pos[v]
+            plt.plot([x1, x2], [y1, y2], color='black', zorder=1)
+    
+    plt.axis('off')
+    chart_path = "/tmp/chart.png"
+    plt.savefig(chart_path, format="png", dpi=300, bbox_inches="tight")
+    plt.close()
+    
+    # Embed chart in HTML for PDF conversion
+    with open(chart_path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode()
     img_html = f"<img src='data:image/png;base64,{encoded}' style='width:100%;'/>"
     html_content = f"""
     <html>
@@ -123,7 +184,6 @@ def generate_chart(file, title):
     </html>
     """
     
-    # Generate PDF.
     wkhtml_path = get_wkhtmltopdf_path()
     if wkhtml_path is None:
         logging.error("wkhtmltopdf not found, creating dummy PDF file.")
@@ -145,6 +205,9 @@ def generate_chart(file, title):
     
     return chart_path, pdf_path
 
+# -------------------------------
+# Template download function
+# -------------------------------
 def download_template():
     template_path = "/tmp/template.csv"
     template_data = "Name,Role,Reporting To\nAlice,CEO,\nBob,Manager,Alice\nCharlie,Engineer,Bob"
@@ -152,6 +215,9 @@ def download_template():
         f.write(template_data)
     return template_path
 
+# -------------------------------
+# Gradio Interface
+# -------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("## Organization Chart Generator")
     gr.Markdown(