Skip to content

Commit 3fe798e

Browse files
authored
Fix to_json method for charts (#68)
* Fix `to_json` method for charts * Add changeset
1 parent 1c4ac6c commit 3fe798e

File tree

3 files changed

+40
-1
lines changed

3 files changed

+40
-1
lines changed

.changeset/smart-turkeys-study.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@e2b/code-interpreter-python': patch
3+
---
4+
5+
Fix json serialization for Result object

python/e2b_code_interpreter/models.py

+8-1
Original file line numberDiff line numberDiff line change
@@ -317,9 +317,16 @@ def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
317317
"""
318318
serialized = []
319319
for result in results:
320-
serialized_dict = {key: result[key] for key in result.formats()}
320+
serialized_dict = {}
321+
for key in result.formats():
322+
if key == "chart":
323+
serialized_dict[key] = result.chart.to_dict()
324+
else:
325+
serialized_dict[key] = result[key]
326+
321327
serialized_dict["text"] = result.text
322328
serialized.append(serialized_dict)
329+
323330
return serialized
324331

325332

python/tests/charts/test_json.py

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import json
2+
3+
from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
4+
5+
code = """
6+
import matplotlib.pyplot as plt
7+
import numpy as np
8+
9+
# Create data
10+
N = 5
11+
x = np.random.rand(N)
12+
y = np.random.rand(N)
13+
14+
plt.xlabel("A")
15+
16+
plt.scatter(x, y, c='blue', label='Dataset')
17+
18+
plt.show()
19+
"""
20+
21+
22+
async def test_scatter_chart(async_sandbox: AsyncSandbox):
23+
result = await async_sandbox.run_code(code)
24+
serialized = result.to_json()
25+
assert isinstance(serialized, str)
26+
27+
assert json.loads(serialized)["results"][0]["chart"]["type"] == "scatter"

0 commit comments

Comments
 (0)