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"