diff --git a/.changeset/smart-turkeys-study.md b/.changeset/smart-turkeys-study.md new file mode 100644 index 00000000..5f346834 --- /dev/null +++ b/.changeset/smart-turkeys-study.md @@ -0,0 +1,5 @@ +--- +'@e2b/code-interpreter-python': patch +--- + +Fix json serialization for Result object diff --git a/python/e2b_code_interpreter/models.py b/python/e2b_code_interpreter/models.py index 04df1a0e..13170a4c 100644 --- a/python/e2b_code_interpreter/models.py +++ b/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 diff --git a/python/tests/charts/test_json.py b/python/tests/charts/test_json.py new file mode 100644 index 00000000..f0dbb189 --- /dev/null +++ b/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"