Update app.py

app.py

@@ -2,24 +2,40 @@ import gradio as gr
 import io
 import contextlib
 import traceback
+import matplotlib.pyplot as plt
 
 def run_python_code(code):
-    output = io.StringIO()
+    stdout = io.StringIO()
+    img_path = "plot.png"
+
+    # Clean up any old plot
+    plt.close()
+
     try:
-        with contextlib.redirect_stdout(output):
-            exec(code, {})  # Run in isolated global scope
-        return output.getvalue()
+        with contextlib.redirect_stdout(stdout):
+            exec(code, {"plt": plt})  # Provide plt to global scope
+
+        # Check if a figure was created
+        figures = [plt.figure(i) for i in plt.get_fignums()]
+        if figures:
+            # Save the latest figure
+            figures[-1].savefig(img_path)
+            return stdout.getvalue(), img_path
+        else:
+            return stdout.getvalue(), None
+
     except Exception:
-        return traceback.format_exc()
+        return traceback.format_exc(), None
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🐍 Python Code Runner (Gradio)")
-    
-    code_input = gr.Code(language="python", label="Write Python Code")
-    output_box = gr.Textbox(label="Output", lines=10)
+    gr.Markdown("## 🐍 Python Code Runner with Plot Support (`matplotlib`, `pandas`)")
     
+    code_input = gr.Code(language="python", label="Enter Python Code")
+    output_text = gr.Textbox(label="Text Output", lines=10)
+    plot_output = gr.Image(label="Generated Plot")
+
     run_button = gr.Button("Run Code")
-    run_button.click(fn=run_python_code, inputs=code_input, outputs=output_box)
+    run_button.click(fn=run_python_code, inputs=code_input, outputs=[output_text, plot_output])
 
 if __name__ == "__main__":
     demo.launch()