Update app.py

app.py

@@ -2,52 +2,50 @@ import gradio as gr
 import io
 import traceback
 import matplotlib.pyplot as plt
-
-from IPython.display import display, HTML
-from contextlib import redirect_stdout
+from IPython.display import display
 from IPython.core.interactiveshell import InteractiveShell
+from IPython.utils.capture import capture_output
 
-# Create a function to capture all kinds of Jupyter-like outputs
 def run_jupyter_style_code(code):
-    # Reset any old figures
+    # Close any old plots
     plt.close('all')
 
-    # Setup for capturing output
+    # Set up an interactive shell instance
     shell = InteractiveShell.instance()
-    stdout = io.StringIO()
+    img_path = None
 
     try:
-        with redirect_stdout(stdout):
-            result = shell.run_cell(code)
-
-        # Collect rich display outputs
-        rich_output = ""
-        if result.result is not None:
-            try:
-                display_out = display(result.result, display_id=False)
-                rich_output = str(result.result)
-            except Exception:
-                rich_output = str(result.result)
-
-        # Show latest matplotlib figure if available
+        # Capture output (both text and rich outputs like HTML, plots)
+        with capture_output(display=True) as captured:
+            exec(code, {"plt": plt, "pd": __import__('pandas'), "display": display})
+
+        # Check for any figures and save the last one as an image
         figures = [plt.figure(i) for i in plt.get_fignums()]
-        img_path = None
         if figures:
             img_path = "plot.png"
             figures[-1].savefig(img_path)
 
-        return stdout.getvalue() + "\n" + rich_output, img_path
+        # Return the captured HTML/text outputs
+        html_output = captured.stdout + ''.join(str(out) for out in captured.outputs)
+        return html_output, img_path
 
     except Exception as e:
+        # Return any errors as string
         return traceback.format_exc(), None
 
+
+# Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("## 🧪 Python Code Runner with Jupyter-style Output")
-
+    
+    # Code input area
     code_input = gr.Code(language="python", label="Enter Python Code")
-    output_text = gr.Textbox(label="Console Output + Rich Result", lines=12)
-    plot_output = gr.Image(label="Plot (if any)")
+    
+    # Outputs (HTML + Plot)
+    output_text = gr.HTML(label="Console Output + Rich Result")
+    plot_output = gr.Image(label="Generated Plot", type="filepath")
 
+    # Run button
     run_button = gr.Button("Run Code")
     run_button.click(fn=run_jupyter_style_code, inputs=code_input, outputs=[output_text, plot_output])