linting / formatting

Author: KRRT7

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -29,7 +29,7 @@ class VerificationType(str, Enum):  # moved from codeflash/verification/codeflas
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-def get_run_tmp_file(file_path: Path) -> Path: # moved from codeflash/code_utils/code_utils.py
+def get_run_tmp_file(file_path: Path) -> Path:  # moved from codeflash/code_utils/code_utils.py
     if not hasattr(get_run_tmp_file, "tmpdir"):
         get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_")
     return Path(get_run_tmp_file.tmpdir.name) / file_path
@@ -181,10 +181,14 @@ def extract_test_context_from_frame() -> tuple[str, str | None, str]:
 
             # Framework integration detection
             if (
-                function_name in ["runTest", "_runTest", "run", "_testMethodName"]
-                or "pytest" in str(frame_obj.f_globals.get("__file__", ""))
-                or "unittest" in str(frame_obj.f_globals.get("__file__", ""))
-            ) and class_name and (class_name.startswith("Test") or "test" in class_name.lower()):
+                (
+                    function_name in ["runTest", "_runTest", "run", "_testMethodName"]
+                    or "pytest" in str(frame_obj.f_globals.get("__file__", ""))
+                    or "unittest" in str(frame_obj.f_globals.get("__file__", ""))
+                )
+                and class_name
+                and (class_name.startswith("Test") or "test" in class_name.lower())
+            ):
                 test_method = function_name
                 if "self" in frame_locals:
                     with contextlib.suppress(AttributeError, TypeError):

codeflash/code_utils/edit_generated_tests.py

@@ -33,8 +33,10 @@ def __init__(
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
         self.context_stack.append(node.name)
         for inner_node in node.body:
-            if isinstance(inner_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            if isinstance(inner_node, ast.FunctionDef):
                 self.visit_FunctionDef(inner_node)
+            elif isinstance(inner_node, ast.AsyncFunctionDef):
+                self.visit_AsyncFunctionDef(inner_node)
         self.context_stack.pop()
         return node
 
@@ -50,14 +52,14 @@ def get_comment(self, match_key: str) -> str:
         return f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
-        return self._process_function_def(node)
+        self._process_function_def_common(node)
+        return node
 
     def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef:
-        return self._process_function_def(node)
+        self._process_function_def_common(node)
+        return node
 
-    def _process_function_def(
-        self, node: ast.FunctionDef | ast.AsyncFunctionDef
-    ) -> ast.FunctionDef | ast.AsyncFunctionDef:
+    def _process_function_def_common(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
         self.context_stack.append(node.name)
         i = len(node.body) - 1
         test_qualified_name = ".".join(self.context_stack)
@@ -83,7 +85,6 @@ def _process_function_def(
                     self.results[line_node.lineno] = self.get_comment(match_key)
             i -= 1
         self.context_stack.pop()
-        return node
 
 
 def get_fn_call_linenos(