From c63f49cb23fcb0fc5f02bf3f02125f2d59e34c86 Mon Sep 17 00:00:00 2001
From: Jakub Novak <jakub@e2b.dev>
Date: Wed, 13 Nov 2024 14:12:35 -0800
Subject: [PATCH] Return json on dict

---
 .changeset/pretty-deers-bow.md        |  5 +++++
 template/startup_scripts/0002_data.py | 16 +++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/pretty-deers-bow.md

diff --git a/.changeset/pretty-deers-bow.md b/.changeset/pretty-deers-bow.md
new file mode 100644
index 0000000..adce63f
--- /dev/null
+++ b/.changeset/pretty-deers-bow.md
@@ -0,0 +1,5 @@
+---
+'@e2b/code-interpreter-template': patch
+---
+
+Return json on dict
diff --git a/template/startup_scripts/0002_data.py b/template/startup_scripts/0002_data.py
index 68e2940..13e916f 100644
--- a/template/startup_scripts/0002_data.py
+++ b/template/startup_scripts/0002_data.py
@@ -1,7 +1,7 @@
 import pandas
 from matplotlib.pyplot import Figure
 import IPython
-from IPython.core.formatters import BaseFormatter
+from IPython.core.formatters import BaseFormatter, JSONFormatter
 from traitlets.traitlets import Unicode, ObjectName
 
 from e2b_charts import chart_figure_to_dict
@@ -48,6 +48,16 @@ def __call__(self, obj):
         return super().__call__(obj)
 
 
+class E2BDictFormatter(JSONFormatter):
+    def __call__(self, obj):
+        # Figure object is for some reason removed on execution of the cell,
+        # so it can't be used in type_printers or with top-level import
+
+        if isinstance(obj, dict):
+            return obj, {"expanded": True}
+        return super().__call__(obj)
+
+
 ip = IPython.get_ipython()
 ip.display_formatter.formatters["e2b/data"] = E2BDataFormatter(
     parent=ip.display_formatter
@@ -55,3 +65,7 @@ def __call__(self, obj):
 ip.display_formatter.formatters["e2b/chart"] = E2BChartFormatter(
     parent=ip.display_formatter
 )
+
+ip.display_formatter.formatters["application/json"] = E2BDictFormatter(
+    parent=ip.display_formatter
+)