Update CodeFileEditAtomicFlow.py

CodeFileEditAtomicFlow.py

@@ -4,34 +4,38 @@ from typing import Dict, Any
 from flows.base_flows.atomic import AtomicFlow
 import os
 
+
 class CodeFileEditAtomicFlow(AtomicFlow):
-    def _write_code_to_temp_file(self, code_str, language_of_code, code_lib_location):
+    def _generate_content(self, code_lib_location, code_str) -> str:
+        content = (
+                "# The below code will be appended to " +
+                code_lib_location + "\n"
+                                    "# Edit the code directly or provide your thoughts down below if you have any suggestions.\n"
+                                    "# Please do provide your feedback in the thoughts section so that JARVIS knows what you "
+                                    "are doing.\n"
+                                    "###########\n"
+                                    "# Code:\n" +
+                code_str +
+                "\n############\n"
+                "# Thoughts:"
+        )
+        return content
+
+    def _generate_temp_file_location(self, code_lib_location):
         directory = os.path.dirname(code_lib_location)
-        if language_of_code.lower() == 'python':
-            temp_file_location = os.path.join(directory, 'temp.py')
-        else:
-            raise NotImplemented
+        ret = os.path.join(directory, 'temp.py')
+        return ret
+
+    def _write_code_content_to_file(self, file_location, content: str):
         try:
-            content = (
-                    "# The below code will be appended to " +
-                    code_lib_location + "\n"
-                    "# Edit the code directly or provide your thoughts down below if you have any suggestions.\n"
-                    "# Please do provide your feedback in the thoughts section so that JARVIS knows what you "
-                    "are doing.\n"
-                    "###########\n"
-                    "# Code:\n" +
-                    code_str +
-                    "\n############\n"
-                    "# Thoughts:"
-            )
-            with open(temp_file_location, "w") as file:
+            with open(file_location, "w") as file:
                 file.write(content)
-            file_written_timestamp = os.path.getmtime(temp_file_location)
+            file_written_timestamp = os.path.getmtime(file_location)
 
-            return True, f"Code written to {temp_file_location}", temp_file_location, file_written_timestamp
+            return True, f"Code written to {file_location}", file_location, file_written_timestamp
 
         except Exception as e:
-            return False, str(e), temp_file_location, 0
+            return False, str(e), file_location, 0
 
     def _check_input(self, input_data: Dict[str, Any]):
         assert "code" in input_data, "code is not passed to CodeFileEditAtomicFlow"
@@ -41,21 +45,33 @@ class CodeFileEditAtomicFlow(AtomicFlow):
         code_lib_loc = input_data["memory_files"]["code_library"]
         assert os.path.exists(code_lib_loc), f"{code_lib_loc} does not exist"
         assert os.path.isfile(code_lib_loc), f"{code_lib_loc} is not a file"
+        language_of_code = input_data["language_of_code"]
+        assert language_of_code.lower() == 'python', "sorry!! only writing python code is supported."
 
     def run(
             self,
             input_data: Dict[str, Any]
     ):
         self._check_input(input_data)
+
+        # ~~~ Generating temp file location and content to write ~~~
         code_str = input_data['code']
-        language_of_code = input_data["language_of_code"]
         code_lib_location = input_data["memory_files"]["code_library"]
-        result, code_editor_output, temp_file_location, file_written_timestamp = self._write_code_to_temp_file(code_str, language_of_code, code_lib_location)
+        content_to_write = self._generate_content(code_lib_location, code_str)
+        file_location_to_write = self._generate_temp_file_location(code_lib_location)
+
+        # ~~~ Calling the writer function ~~~
+        result, code_editor_output, temp_file_location, file_written_timestamp = self._write_code_content_to_file(
+            file_location_to_write, content_to_write)
+
+        # ~~~ Opening up the file for the user to see ~~~
         if result:
             try:
                 subprocess.run(["code", temp_file_location], timeout=10)
             except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                 webbrowser.open(temp_file_location)
+
+        # ~~~ Generating return variables ~~~
         response = {}
         response["code_editor_output"] = code_editor_output
         response["temp_code_file_location"] = temp_file_location