Update tools/plot_generator.py

tools/plot_generator.py

@@ -1,7 +1,9 @@
+# tools/plot_generator.py
 import os
 import tempfile
 import pandas as pd
 import plotly.graph_objects as go
+from typing import Tuple, Union
 
 
 def plot_metric_tool(
@@ -12,30 +14,28 @@ def plot_metric_tool(
     title: str = None,
     line_width: int = 2,
     marker_size: int = 6
-):
+) -> Union[Tuple[go.Figure, str], str]:
     """
-    Load CSV or Excel file, parse a time series metric, and return an interactive Plotly Figure.
-    Also saves a high-resolution PNG to a temp directory for static embedding.
+    Load CSV or Excel file, parse a time series metric, and return an interactive Plotly Figure
+    plus a high-res PNG file path for static embedding.
 
     Returns:
-      fig (go.Figure) or error string starting with '❌'.
+      - (fig, img_path) on success
+      - error string starting with '❌' on failure
     """
-    # 0) Load data
+    # Load data
     ext = os.path.splitext(file_path)[1].lower()
     try:
-        if ext in ('.xls', '.xlsx'):
-            df = pd.read_excel(file_path)
-        else:
-            df = pd.read_csv(file_path)
+        df = pd.read_excel(file_path) if ext in ('.xls', '.xlsx') else pd.read_csv(file_path)
     except Exception as exc:
         return f"❌ Failed to load file: {exc}"
 
-    # 1) Validate columns
+    # Validate columns
     missing = [c for c in (date_col, value_col) if c not in df.columns]
     if missing:
         return f"❌ Missing column(s): {', '.join(missing)}"
 
-    # 2) Parse date and ensure numeric values
+    # Parse and clean
     try:
         df[date_col] = pd.to_datetime(df[date_col], errors='coerce')
     except Exception:
@@ -45,7 +45,7 @@ def plot_metric_tool(
     if df.empty:
         return f"❌ No valid data after cleaning '{date_col}'/'{value_col}'"
 
-    # 3) Sort and aggregate duplicates
+    # Aggregate duplicates and sort
     df = (
         df[[date_col, value_col]]
         .groupby(date_col, as_index=True)
@@ -53,7 +53,7 @@ def plot_metric_tool(
         .sort_index()
     )
 
-    # 4) Create Plotly figure
+    # Build figure
     fig = go.Figure(
         data=[
             go.Scatter(
@@ -62,30 +62,26 @@ def plot_metric_tool(
                 mode='lines+markers',
                 line=dict(width=line_width),
                 marker=dict(size=marker_size),
-                name=value_col
+                name=value_col,
             )
         ]
     )
-    plot_title = title or f"{value_col} Trend"
     fig.update_layout(
-        title=plot_title,
+        title=title or f"{value_col} Trend",
         xaxis_title=date_col,
         yaxis_title=value_col,
         template='plotly_dark',
         hovermode='x unified'
     )
 
-    # 5) Save static PNG
+    # Save PNG
     os.makedirs(output_dir, exist_ok=True)
-    tmpfile = tempfile.NamedTemporaryFile(
-        suffix='.png', prefix='trend_', dir=output_dir, delete=False
-    )
-    img_path = tmpfile.name
-    tmpfile.close()
+    tmp = tempfile.NamedTemporaryFile(suffix='.png', prefix='trend_', dir=output_dir, delete=False)
+    img_path = tmp.name
+    tmp.close()
     try:
         fig.write_image(img_path, scale=2)
     except Exception as exc:
         return f"❌ Failed saving image: {exc}"
 
-    # 6) Return figure and path for embedding
     return fig, img_path