Fix to_json method for charts

Author: jakubno

python/e2b_code_interpreter/models.py

@@ -317,9 +317,16 @@ def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
     """
     serialized = []
     for result in results:
-        serialized_dict = {key: result[key] for key in result.formats()}
+        serialized_dict = {}
+        for key in result.formats():
+            if key == "chart":
+                serialized_dict[key] = result.chart.to_dict()
+            else:
+                serialized_dict[key] = result[key]
+
         serialized_dict["text"] = result.text
         serialized.append(serialized_dict)
+
     return serialized
 
 

python/tests/charts/test_json.py

@@ -0,0 +1,27 @@
+import json
+
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+code = """
+import matplotlib.pyplot as plt
+import numpy as np
+
+# Create data
+N = 5
+x = np.random.rand(N)
+y = np.random.rand(N)
+
+plt.xlabel("A")
+
+plt.scatter(x, y, c='blue', label='Dataset')
+
+plt.show()
+"""
+
+
+async def test_scatter_chart(async_sandbox: AsyncSandbox):
+    result = await async_sandbox.run_code(code)
+    serialized = result.to_json()
+    assert isinstance(serialized, str)
+
+    assert json.loads(serialized)["results"][0]["chart"]["type"] == "scatter"