| import os | |
| from typing import Dict, Any | |
| from flow_modules.aiflows.CodeFileEditFlowModule import CodeFileEditAtomicFlow | |
| class TestCodeFileEditAtomicFlow(CodeFileEditAtomicFlow): | |
| """Refer to: https://huggingface.co/aiflows/CodeFileEditFlowModule/tree/main | |
| *Input Interface*: | |
| - `code`: str | |
| - `memory_files`: Dict[str, str] | |
| *Output Interface*: | |
| - `code_editor_output`: str, the code editor output | |
| - `temp_code_file_location`: str, the location of the temporary code file | |
| """ | |
| def _generate_import_statement(self, code_lib_location): | |
| """ | |
| Generate the import statement for the code library. | |
| :param code_lib_location: the location of the code library | |
| :type code_lib_location: str | |
| :return: the import statement | |
| :rtype: str | |
| """ | |
| module_dir = os.path.dirname(code_lib_location) | |
| module_name = os.path.splitext(os.path.basename(code_lib_location))[0] | |
| import_code = ( | |
| f"import sys\n" | |
| f"sys.path.insert(0, '{module_dir}')\n" | |
| f"from {module_name} import *\n" | |
| ) | |
| return import_code | |
| def _generate_content(self, code_lib_location, code_str) -> str: | |
| """ | |
| Generate the content of the temporary code file. | |
| :param code_lib_location: the location of the code library | |
| :type code_lib_location: str | |
| :param code_str: the code string | |
| :type code_str: str | |
| :return: the content of the temporary code file | |
| :rtype: str | |
| """ | |
| import_code_lib_str = self._generate_import_statement(code_lib_location) | |
| content = ( | |
| "# Don't touch this import statement \n" | |
| + import_code_lib_str + "\n" | |
| "# Here is the code just generated \n" + | |
| code_str + "\n" | |
| "# Below, please provide code to test it.\n" | |
| "# The simplest form could be just calling it with appropriate parameters. \n" | |
| "# You could also assert the output, anyway, the test results will be informed to JARVIS. \n" | |
| "# If you do not write anything, JARVIS just checks if the syntax is alright. \n" | |
| "###########\n" | |
| "# Test Code:\n" + | |
| "\n############\n" | |
| ) | |
| return content | |
| def _generate_temp_file_location(self, code_lib_location): | |
| """ | |
| Generate the location of the temporary code file. | |
| :param code_lib_location: the location of the code library | |
| :type code_lib_location: str | |
| :return: the location of the temporary code file | |
| :rtype: str | |
| """ | |
| directory = os.path.dirname(code_lib_location) | |
| ret = os.path.join(directory, 'temp_tests.py') | |
| return ret | |
| def _check_input(self, input_data: Dict[str, Any]): | |
| """ | |
| Check if the input data is valid. | |
| :param input_data: the input data | |
| :type input_data: Dict[str, Any] | |
| :raises AssertionError: if code or memory_files is not passed to TestCodeFileEditAtomicFlow | |
| """ | |
| assert "code" in input_data, "code is not passed to TestCodeFileEditAtomicFlow" | |
| assert "memory_files" in input_data, "memory_files is not passed to TestCodeFileEditAtomicFlow" | |