Update app.py

app.py

@@ -49,18 +49,32 @@ async def run_code(code: Code):
 
     file_path = "graph.pdf"
 
-    if "plt.subplots(" in code.code:
-        print("SUBPLOTS DETECTED")
-        exec(code.code + "\nfig.savefig(file_path)\nplt.close()")
-    else:
-        print("NO SUBPLOTS")
-        exec(code.code + "\nplt.savefig(file_path)\nplt.close()")
+    # if "plt.subplots(" in code.code:
+    #     print("SUBPLOTS DETECTED")
+    #     exec(code.code + "\nfig.savefig(file_path)\nplt.close()")
+    # else:
+    #     print("NO SUBPLOTS")
+    #     exec(code.code + "\nplt.savefig(file_path)\nplt.close()")
 
-    plt.close()
+    # plt.close()
     # plt.savefig(file_path)
     # plt.close()
 
-    # Return image as response
+    # Get all open figure numbers
+    figures = [plt.figure(i) for i in plt.get_fignums()]
+
+    
+    # Save all figures in a single PDF
+    if figures:
+        figures[0].savefig("graph.pdf", format="pdf")
+        for fig in figures[1:]:
+            fig.savefig("graph.pdf", format="pdf", bbox_inches="tight", append=True)
+        
+        plt.close("all")  # Close all figures
+    else:
+        print("No open figures found.")
+
+    
     # return Response(content=img_buffer.getvalue(), media_type="image/png")
     # return FileResponse(file_path, media_type="image/png")
     return FileResponse(file_path, media_type="application/pdf", filename="graph.pdf")