refactor

Author: KRRT7

codeflash/optimization/optimizer.py

@@ -190,9 +190,23 @@ def create_function_optimizer(
         function_to_optimize_source_code: str | None = "",
         function_benchmark_timings: dict[str, dict[BenchmarkKey, float]] | None = None,
         total_benchmark_timings: dict[BenchmarkKey, float] | None = None,
-    ) -> FunctionOptimizer:
+        original_module_ast: ast.Module | None = None,
+        original_module_path: Path | None = None,
+    ) -> FunctionOptimizer | None:
+        from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
         from codeflash.optimization.function_optimizer import FunctionOptimizer
 
+        if function_to_optimize_ast is None and original_module_ast is not None:
+            function_to_optimize_ast = get_first_top_level_function_or_method_ast(
+                function_to_optimize.function_name, function_to_optimize.parents, original_module_ast
+            )
+            if function_to_optimize_ast is None:
+                logger.info(
+                    f"Function {function_to_optimize.qualified_name} not found in {original_module_path}.\n"
+                    f"Skipping optimization."
+                )
+                return None
+
         qualified_name_w_module = function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
 
         function_specific_timings = None
@@ -220,7 +234,6 @@ def create_function_optimizer(
     def run(self) -> None:
         from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint
         from codeflash.code_utils.code_utils import cleanup_paths
-        from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast
 
         ph("cli-optimize-run-start")
         logger.info("Running optimizer.")
@@ -266,26 +279,20 @@ def run(self) -> None:
                         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
-                        )
-                    ):
-                        logger.info(
-                            f"Function {function_to_optimize.qualified_name} not found in {original_module_path}.\n"
-                            f"Skipping optimization."
-                        )
-                        continue
 
                     function_optimizer = self.create_function_optimizer(
                         function_to_optimize,
-                        function_to_optimize_ast,
-                        function_to_tests,
-                        validated_original_code[original_module_path].source_code,
-                        function_benchmark_timings,
-                        total_benchmark_timings,
+                        function_to_tests=function_to_tests,
+                        function_to_optimize_source_code=validated_original_code[original_module_path].source_code,
+                        function_benchmark_timings=function_benchmark_timings,
+                        total_benchmark_timings=total_benchmark_timings,
+                        original_module_ast=original_module_ast,
+                        original_module_path=original_module_path,
                     )
 
+                    if function_optimizer is None:
+                        continue
+
                     best_optimization = function_optimizer.optimize_function()
                     if self.functions_checkpoint:
                         self.functions_checkpoint.add_function_to_checkpoint(