update: 添加nb支持

tests/metagpt/utils/test_save_code.py

@@ -2,8 +2,13 @@
 # @Date    : 12/12/2023 4:17 PM
 # @Author  : stellahong (stellahong@fuzhi.ai)
 # @Desc    :
+import pytest
 import os
 import json
+import nbformat
+
+from metagpt.actions.write_analysis_code import WriteCodeByGenerate
+from metagpt.actions.execute_code import ExecutePyCode
 
 from metagpt.utils.save_code import save_code_file, DATA_PATH
 
@@ -21,6 +26,7 @@ def test_save_code_file_python():
         content = fp.read()
     assert "print('Hello, World!')" in content, "File content does not match"
 
+
 def test_save_code_file_json():
     save_code_file("example_json", "print('Hello, JSON!')", file_format="json")
     file_path = DATA_PATH / "output" / "example_json" / "code.json"
@@ -28,3 +34,23 @@ def test_save_code_file_json():
         data = json.load(fp)
     assert "code" in data, "JSON key 'code' is missing"
     assert data["code"] == "print('Hello, JSON!')", "JSON content does not match"
+
+
+
+@pytest.mark.asyncio
+async def test_save_code_file_notebook():
+    code = await WriteCodeByGenerate().run(
+        context="basic python, hello world", plan="", code_steps="", temperature=0.0
+    )
+    executor = ExecutePyCode()
+    await executor.run(code)
+    # Save as a Notebook file
+    save_code_file("example_nb", executor.nb, file_format="ipynb")
+    file_path = DATA_PATH / "output" / "example_nb" / "code.ipynb"
+    assert os.path.exists(file_path), f"Notebook file does not exist: {file_path}"
+
+    # Additional checks specific to notebook format
+    notebook = nbformat.read(file_path, as_version=4)
+    assert len(notebook.cells) > 0, "Notebook should have at least one cell"
+    first_cell_source = notebook.cells[0].source
+    assert "print('Hello, World!')" in first_cell_source, "Notebook cell content does not match"