Merge pull request #33 from codeflash-ai/skip-async-dependent

Skip async dependent functions

Author: KRRT7

codeflash/code_utils/code_utils.py

@@ -108,6 +108,14 @@ def validate_python_code(code: str) -> str:
     return code
 
 
+def has_any_async_functions(code: str) -> bool:
+    try:
+        module = ast.parse(code)
+    except SyntaxError:
+        return False
+    return any(isinstance(node, ast.AsyncFunctionDef) for node in ast.walk(module))
+
+
 def cleanup_paths(paths: list[Path]) -> None:
     for path in paths:
         path.unlink(missing_ok=True)

codeflash/optimization/function_optimizer.py

@@ -28,6 +28,7 @@
     file_name_from_test_module_name,
     get_run_tmp_file,
     module_name_from_file_path,
+    has_any_async_functions,
 )
 from codeflash.code_utils.config_consts import (
     INDIVIDUAL_TESTCASE_TIMEOUT,
@@ -133,8 +134,8 @@ def optimize_function(self) -> Result[BestOptimization, str]:
             with helper_function_path.open(encoding="utf8") as f:
                 helper_code = f.read()
                 original_helper_code[helper_function_path] = helper_code
-
-        logger.info("Code to be optimized:")
+        if has_any_async_functions(code_context.code_to_optimize_with_helpers):
+            return Failure("Codeflash does not support async functions in the code to optimize.")
         code_print(code_context.read_writable_code)
 
         for module_abspath, helper_code_source in original_helper_code.items():

codeflash/optimization/optimizer.py

@@ -139,6 +139,7 @@ def run(self) -> None:
                     validated_original_code[analysis.file_path] = ValidCode(
                         source_code=callee_original_code, normalized_code=normalized_callee_original_code
                     )
+
                 if has_syntax_error:
                     continue
 
@@ -148,7 +149,7 @@ def run(self) -> None:
                         f"Optimizing function {function_iterator_count} of {num_optimizable_functions}: "
                         f"{function_to_optimize.qualified_name}"
                     )
-
+                    console.rule()
                     if not (
                         function_to_optimize_ast := get_first_top_level_function_or_method_ast(
                             function_to_optimize.function_name, function_to_optimize.parents, original_module_ast
@@ -159,7 +160,6 @@ def run(self) -> None:
                             f"Skipping optimization."
                         )
                         continue
-
                     function_optimizer = self.create_function_optimizer(
                         function_to_optimize,
                         function_to_optimize_ast,

tests/test_code_utils.py

@@ -17,6 +17,7 @@
     is_class_defined_in_file,
     module_name_from_file_path,
     path_belongs_to_site_packages,
+    has_any_async_functions,
 )
 from codeflash.code_utils.concolic_utils import clean_concolic_tests
 from codeflash.code_utils.coverage_utils import generate_candidates, prepare_coverage_files
@@ -441,3 +442,27 @@ def test_Grammar_copy():
     Grammar.copy(Grammar())
 """
     assert cleaned_code == expected_cleaned_code.strip()
+
+
+def test_has_any_async_functions_with_async_code() -> None:
+    code = """
+def normal_function():
+    pass
+
+async def async_function():
+    pass
+"""
+    result = has_any_async_functions(code)
+    assert result is True
+
+
+def test_has_any_async_functions_without_async_code() -> None:
+    code = """
+def normal_function():
+    pass
+
+def another_function():
+    pass
+"""
+    result = has_any_async_functions(code)
+    assert result is False