gate async behind --async

Author: KRRT7

codeflash/cli_cmds/cli.py

@@ -96,6 +96,9 @@ def parse_args() -> Namespace:
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
     parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
+    parser.add_argument(
+        "--async", default=False, action="store_true", help="Enable optimization of async functions. By default, async functions are excluded from optimization."
+    )
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]

codeflash/discovery/functions_to_optimize.py

@@ -179,6 +179,7 @@ def get_functions_to_optimize(
     project_root: Path,
     module_root: Path,
     previous_checkpoint_functions: dict[str, dict[str, str]] | None = None,
+    enable_async: bool = False,
 ) -> tuple[dict[Path, list[FunctionToOptimize]], int, Path | None]:
     assert sum([bool(optimize_all), bool(replay_test), bool(file)]) <= 1, (
         "Only one of optimize_all, replay_test, or file should be provided"
@@ -232,7 +233,7 @@ def get_functions_to_optimize(
             ph("cli-optimizing-git-diff")
             functions = get_functions_within_git_diff(uncommitted_changes=False)
         filtered_modified_functions, functions_count = filter_functions(
-            functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
+            functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions, enable_async=enable_async
         )
 
         logger.info(f"!lsp|Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
@@ -568,6 +569,7 @@ def filter_functions(
     module_root: Path,
     previous_checkpoint_functions: dict[Path, dict[str, Any]] | None = None,
     disable_logs: bool = False,  # noqa: FBT001, FBT002
+    enable_async: bool = False,
 ) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
     filtered_modified_functions: dict[str, list[FunctionToOptimize]] = {}
     blocklist_funcs = get_blocklisted_functions()
@@ -587,6 +589,7 @@ def filter_functions(
     submodule_ignored_paths_count: int = 0
     blocklist_funcs_removed_count: int = 0
     previous_checkpoint_functions_removed_count: int = 0
+    async_functions_removed_count: int = 0
     tests_root_str = str(tests_root)
     module_root_str = str(module_root)
 
@@ -642,6 +645,15 @@ def filter_functions(
                 functions_tmp.append(function)
             _functions = functions_tmp
 
+        if not enable_async:
+            functions_tmp = []
+            for function in _functions:
+                if function.is_async:
+                    async_functions_removed_count += 1
+                    continue
+                functions_tmp.append(function)
+            _functions = functions_tmp
+
         filtered_modified_functions[file_path] = _functions
         functions_count += len(_functions)
 
@@ -655,6 +667,7 @@ def filter_functions(
             "Files from ignored submodules": (submodule_ignored_paths_count, "bright_black"),
             "Blocklisted functions removed": (blocklist_funcs_removed_count, "bright_red"),
             "Functions skipped from checkpoint": (previous_checkpoint_functions_removed_count, "green"),
+            "Async functions removed": (async_functions_removed_count, "bright_magenta"),
         }
         tree = Tree(Text("Ignored functions and files", style="bold"))
         for label, (count, color) in log_info.items():

codeflash/optimization/optimizer.py

@@ -134,6 +134,7 @@ def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize]
             project_root=self.args.project_root,
             module_root=self.args.module_root,
             previous_checkpoint_functions=self.args.previous_checkpoint_functions,
+            enable_async=getattr(self.args, "async", False),
         )
 
     def create_function_optimizer(

tests/test_async_function_discovery.py

@@ -236,6 +236,7 @@ def sync_method(self):
         ignore_paths=[],
         project_root=file_path.parent,
         module_root=file_path.parent,
+        enable_async=True,
     )
 
     assert functions_count == 4
@@ -249,6 +250,56 @@ def sync_method(self):
     assert "async_func_two" not in function_names
 
 
+def test_no_async_functions_finding(temp_dir):
+    mixed_code = """
+async def async_func_one():
+    return await operation_one()
+
+def sync_func_one():
+    return operation_one()
+
+async def async_func_two():
+    print("no return")
+
+class MixedClass:
+    async def async_method(self):
+        return await self.operation()
+    
+    def sync_method(self):
+        return self.operation()
+"""
+    
+    file_path = temp_dir / "test_file.py"
+    file_path.write_text(mixed_code)
+    
+    test_config = TestConfig(
+        tests_root="tests",
+        project_root_path=".",
+        test_framework="pytest",
+        tests_project_rootdir=Path()
+    )
+    
+    functions, functions_count, _ = get_functions_to_optimize(
+        optimize_all=None,
+        replay_test=None,
+        file=file_path,
+        only_get_this_function=None,
+        test_cfg=test_config,
+        ignore_paths=[],
+        project_root=file_path.parent,
+        module_root=file_path.parent,
+        enable_async=False,
+    )
+    
+    assert functions_count == 2
+    
+    function_names = [fn.function_name for fn in functions[file_path]]
+    assert "sync_func_one" in function_names
+    assert "sync_method" in function_names
+    assert "async_func_one" not in function_names
+    assert "async_method" not in function_names
+
+
 def test_async_function_parents(temp_dir):
     complex_structure = """
 class OuterClass: