MetaGPT

1
# -*- coding: utf-8 -*-
2
# @Date    : 12/12/2023 4:14 PM
3
# @Author  : stellahong (stellahong@fuzhi.ai)
4
# @Desc    :
5
import os
6

7
import nbformat
8

9
from metagpt.const import DATA_PATH
10
from metagpt.utils.common import write_json_file
11

12

13
def save_code_file(name: str, code_context: str, file_format: str = "py") -> None:
14
    """
15
    Save code files to a specified path.
16

17
    Args:
18
    - name (str): The name of the folder to save the files.
19
    - code_context (str): The code content.
20
    - file_format (str, optional): The file format. Supports 'py' (Python file), 'json' (JSON file), and 'ipynb' (Jupyter Notebook file). Default is 'py'.
21

22

23
    Returns:
24
    - None
25
    """
26
    # Create the folder path if it doesn't exist
27
    os.makedirs(name=DATA_PATH / "output" / f"{name}", exist_ok=True)
28

29
    # Choose to save as a Python file or a JSON file based on the file format
30
    file_path = DATA_PATH / "output" / f"{name}/code.{file_format}"
31
    if file_format == "py":
32
        file_path.write_text(code_context + "\n\n", encoding="utf-8")
33
    elif file_format == "json":
34
        # Parse the code content as JSON and save
35
        data = {"code": code_context}
36
        write_json_file(file_path, data, encoding="utf-8", indent=2)
37
    elif file_format == "ipynb":
38
        nbformat.write(code_context, file_path)
39
    else:
40
        raise ValueError("Unsupported file format. Please choose 'py', 'json', or 'ipynb'.")
41