feat: Add plot export functinality

app.py

@@ -50,6 +50,7 @@ from visualizations import (
     create_distribution_plot,
     create_scatter_plot,
     create_time_series_plot,
+    figure_to_png_bytes,
 )
 
 
@@ -79,6 +80,7 @@ DEFAULT_STATE = {
     "column_types": None,
     "filter_meta": None,
     "source_name": None,
+    "current_figure": None,
 }
 
 
@@ -360,13 +362,14 @@ def _generate_chart(
     scatter_x: Optional[str],
     scatter_y: Optional[str],
     scatter_color: Optional[str],
-) -> Tuple[Any, Any, str]:
+) -> Tuple[Dict[str, Any], Any, str]:
     """Create a visualization based on user selections."""
     state = _ensure_state(state)
     try:
         df = _current_dataframe(state, filtered=True)
     except ValueError as exc:
-        return None, None, f"⚠️ {exc}"
+        state["current_figure"] = None
+        return state, None, f"⚠️ {exc}"
 
     try:
         if chart_type == "Time Series":
@@ -390,9 +393,11 @@ def _generate_chart(
         else:
             raise ValueError("Unsupported chart type.")
     except ValueError as exc:
-        return None, None, f"⚠️ {exc}"
+        state["current_figure"] = None
+        return state, None, f"⚠️ {exc}"
 
-    return fig, "Visualization generated."
+    state["current_figure"] = fig
+    return state, fig, "✅ Visualization generated. Use 'Export Chart' to download."
 
 
 def _download_filtered(state) -> str:
@@ -408,8 +413,10 @@ def _download_filtered(state) -> str:
     return temp.name
 
 
-def _download_chart(fig: Any) -> str:
+def _download_chart(state) -> str:
     """Export the most recent chart to PNG."""
+    state = _ensure_state(state)
+    fig = state.get("current_figure")
     if fig is None:
         raise ValueError("Generate a visualization before exporting.")
     buffer = figure_to_png_bytes(fig)
@@ -551,8 +558,11 @@ def create_dashboard():
                 scatter_y_column = gr.Dropdown(label="Y Axis", choices=[])
                 scatter_color_column = gr.Dropdown(label="Color (optional)", choices=[])
 
-            generate_chart_button = gr.Button("Generate Visualization", variant="primary")
+            with gr.Row():
+                generate_chart_button = gr.Button("Generate Visualization", variant="primary")
+                export_chart_button = gr.Button("Export Chart (PNG)", variant="secondary")
             chart_output = gr.Plot(label="Visualization")
+            export_chart_file = gr.File(label="Exported Chart", interactive=False)
 
 
         with gr.Tab("Insights"):
@@ -744,7 +754,13 @@ def create_dashboard():
                 scatter_y_column,
                 scatter_color_column,
             ],
-            outputs=[chart_output, viz_status],
+            outputs=[dataset_state, chart_output, viz_status],
+        )
+
+        export_chart_button.click(
+            fn=_download_chart,
+            inputs=[dataset_state],
+            outputs=[export_chart_file],
         )