From 0231cfdcc750f3366c3eee16fc776581f67cbaf6 Mon Sep 17 00:00:00 2001
From: stellahsr <stellahsr@126.com>
Date: Tue, 12 Dec 2023 16:23:56 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=BB=A3=E7=A0=81=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=E8=BE=93=E5=87=BA=E4=BF=9D=E5=AD=98=EF=BC=8C=E5=88=9B?=
 =?UTF-8?q?=E5=BB=BA=E9=A1=B9=E7=9B=AE=E6=96=87=E4=BB=B6=E5=A4=B9=EF=BC=8C?=
 =?UTF-8?q?=E4=BD=BF=E7=94=A8=E9=A1=B9=E7=9B=AE=E6=96=87=E4=BB=B6=E5=A4=B9?=
 =?UTF-8?q?=E9=9A=94=E7=A6=BB=20=E5=AE=8C=E6=95=B4=E4=BB=A3=E7=A0=81?=
 =?UTF-8?q?=E4=BF=9D=E5=AD=98=E5=89=8D=EF=BC=8C=E5=8F=AF=E8=80=83=E8=99=91?=
 =?UTF-8?q?=E6=8B=BC=E6=8E=A5=E5=85=A8=E9=87=8F=E4=BB=A3=E7=A0=81=E5=86=8D?=
 =?UTF-8?q?=E8=BE=93=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 metagpt/utils/save_code.py            | 40 +++++++++++++++++++++++++++
 tests/metagpt/utils/test_save_code.py | 30 ++++++++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100644 metagpt/utils/save_code.py
 create mode 100644 tests/metagpt/utils/test_save_code.py

diff --git a/metagpt/utils/save_code.py b/metagpt/utils/save_code.py
new file mode 100644
index 000000000..b0720a5cf
--- /dev/null
+++ b/metagpt/utils/save_code.py
@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# @Date    : 12/12/2023 4:14 PM
+# @Author  : stellahong (stellahong@fuzhi.ai)
+# @Desc    :
+import os
+import json
+
+from metagpt.const import DATA_PATH
+
+def save_code_file(name: str, code_context: str, file_format: str = "py") -> None:
+    """
+    Save code files to a specified path.
+
+    Args:
+    - name (str): The name of the folder to save the files.
+    - code_context (str): The code content.
+    - file_format (str, optional): The file format, supports 'py' (Python file) and 'json' (JSON file). Default is 'py'.
+
+    Returns:
+    - None
+    """
+    # Create the folder path if it doesn't exist
+    os.makedirs(name=DATA_PATH / "output" / f"{name}", exist_ok=True)
+
+    # Choose to save as a Python file or a JSON file based on the file format
+    file_path = DATA_PATH / "output" / f"{name}/code.{file_format}"
+    if file_format == "py":
+        with open(file_path, "w", encoding="utf-8") as fp:
+            fp.write(code_context + "\n\n")
+    elif file_format == "json":
+        # Parse the code content as JSON and save
+        data = {"code": code_context}
+        with open(file_path, "w", encoding="utf-8") as fp:
+            json.dump(data, fp, indent=2)
+    else:
+        raise ValueError("Unsupported file format. Please choose 'py' or 'json'.")
+
+
+
+
diff --git a/tests/metagpt/utils/test_save_code.py b/tests/metagpt/utils/test_save_code.py
new file mode 100644
index 000000000..33addb2bf
--- /dev/null
+++ b/tests/metagpt/utils/test_save_code.py
@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# @Date    : 12/12/2023 4:17 PM
+# @Author  : stellahong (stellahong@fuzhi.ai)
+# @Desc    :
+import os
+import json
+
+from metagpt.utils.save_code import save_code_file, DATA_PATH
+
+
+def test_save_code_file_python():
+    save_code_file("example", "print('Hello, World!')")
+    file_path = DATA_PATH / "output" / "example" / "code.py"
+    assert os.path.exists(file_path), f"File does not exist: {file_path}"
+
+
+def test_save_code_file_python():
+    save_code_file("example", "print('Hello, World!')")
+    file_path = DATA_PATH / "output" / "example" / "code.py"
+    with open(file_path, "r", encoding="utf-8") as fp:
+        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"
+    with open(file_path, "r", encoding="utf-8") as fp:
+        data = json.load(fp)
+    assert "code" in data, "JSON key 'code' is missing"
+    assert data["code"] == "print('Hello, JSON!')", "JSON content does not match"