restore from other branch

Author: KRRT7

codeflash/code_utils/static_analysis.py

@@ -128,13 +128,19 @@ def get_first_top_level_object_def_ast(
 
 def get_first_top_level_function_or_method_ast(
     function_name: str, parents: list[FunctionParent], node: ast.AST
-) -> ast.FunctionDef | None:
+) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
     if not parents:
-        return get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node)
+        result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node)
+        if result is not None:
+            return result
+        return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, node)
     if parents[0].type == "ClassDef" and (
         class_node := get_first_top_level_object_def_ast(parents[0].name, ast.ClassDef, node)
     ):
-        return get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
+        result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
+        if result is not None:
+            return result
+        return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, class_node)
     return None
 
 

codeflash/optimization/optimizer.py

@@ -139,7 +139,7 @@ def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize]
     def create_function_optimizer(
         self,
         function_to_optimize: FunctionToOptimize,
-        function_to_optimize_ast: ast.FunctionDef | None = None,
+        function_to_optimize_ast: ast.FunctionDef | ast.AsyncFunctionDef | None = None,
         function_to_tests: dict[str, set[FunctionCalledInTest]] | None = None,
         function_to_optimize_source_code: str | None = "",
         function_benchmark_timings: dict[str, dict[BenchmarkKey, float]] | None = None,