Update app.py

app.py

@@ -1,41 +1,55 @@
 import gradio as gr
 import io
-import contextlib
 import traceback
 import matplotlib.pyplot as plt
 
-def run_python_code(code):
-    stdout = io.StringIO()
-    img_path = "plot.png"
+from IPython.display import display, HTML
+from contextlib import redirect_stdout
+from IPython.core.interactiveshell import InteractiveShell
 
-    # Clean up any old plot
-    plt.close()
+# Create a function to capture all kinds of Jupyter-like outputs
+def run_jupyter_style_code(code):
+    # Reset any old figures
+    plt.close('all')
 
-    try:
-        with contextlib.redirect_stdout(stdout):
-            exec(code, {"plt": plt})  # Provide plt to global scope
+    # Setup for capturing output
+    shell = InteractiveShell.instance()
+    stdout = io.StringIO()
 
-        # Check if a figure was created
+    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
         figures = [plt.figure(i) for i in plt.get_fignums()]
+        img_path = None
         if figures:
-            # Save the latest figure
+            img_path = "plot.png"
             figures[-1].savefig(img_path)
-            return stdout.getvalue(), img_path
-        else:
-            return stdout.getvalue(), None
 
-    except Exception:
+        return stdout.getvalue() + "\n" + rich_output, img_path
+
+    except Exception as e:
         return traceback.format_exc(), None
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🐍 Python Code Runner with Plot Support (`matplotlib`, `pandas`)")
-    
+    gr.Markdown("## 🧪 Python Code Runner with Jupyter-style Output")
+
     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")
+    output_text = gr.Textbox(label="Console Output + Rich Result", lines=12)
+    plot_output = gr.Image(label="Plot (if any)")
 
     run_button = gr.Button("Run Code")
-    run_button.click(fn=run_python_code, inputs=code_input, outputs=[output_text, plot_output])
+    run_button.click(fn=run_jupyter_style_code, inputs=code_input, outputs=[output_text, plot_output])
 
 if __name__ == "__main__":
     demo.launch()