shifted cleanup

Author: alvin-r

codeflash/optimization/function_optimizer.py

@@ -545,10 +545,6 @@ def replace_function_and_helpers_with_optimized_code(
         for helper_function in code_context.helper_functions:
             if helper_function.jedi_definition.type != "class":
                 read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
-        logger.debug("Read writable functions by file path:")
-        logger.debug(read_writable_functions_by_file_path)
-        logger.debug("Optimized code:")
-        logger.debug(optimized_code)
         for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
             did_update |= replace_function_definitions_in_module(
                 function_names=list(qualified_names),
@@ -834,6 +830,7 @@ def establish_original_code_baseline(
 
             coverage_results = None
             # Instrument codeflash capture
+            original_fto_code = Path(self.function_to_optimize.file_path).read_text("utf-8")
             try:
                 instrument_codeflash_capture(
                     self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
@@ -850,15 +847,18 @@ def establish_original_code_baseline(
             finally:
                 # Remove codeflash capture
                 self.write_code_and_helpers(
-                    self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+                    original_fto_code, original_helper_code, self.function_to_optimize.file_path
                 )
             if not behavioral_results:
                 logger.warning(
                     f"Couldn't run any tests for original function {self.function_to_optimize.function_name}. SKIPPING OPTIMIZING THIS FUNCTION."
                 )
                 console.rule()
                 return Failure("Failed to establish a baseline for the original code - bevhavioral tests failed.")
-
+            if not coverage_critic(
+            coverage_results, self.args.test_framework
+                ):
+                return Failure("The threshold for test coverage was not met.")
             if test_framework == "pytest":
                 benchmarking_results, _ = self.run_and_parse_tests(
                     testing_type=TestingMode.PERFORMANCE,

codeflash/optimization/optimizer.py

@@ -71,7 +71,7 @@ def run(self) -> None:
         console.rule()
         if not env_utils.ensure_codeflash_api_key():
             return
-
+        function_optimizer = None
         file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]]
         num_optimizable_functions: int
         (file_to_funcs_to_optimize, num_optimizable_functions) = get_functions_to_optimize(
@@ -166,18 +166,7 @@ def run(self) -> None:
                     function_optimizer = self.create_function_optimizer(
                         function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code
                     )
-                    try:
-                        best_optimization = function_optimizer.optimize_function()
-                    finally:
-                        for test_file in function_optimizer.test_files.get_by_type(TestType.GENERATED_REGRESSION).test_files:
-                            test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
-                            test_file.benchmarking_file_path.unlink(missing_ok=True)
-                        for test_file in function_optimizer.test_files.get_by_type(TestType.EXISTING_UNIT_TEST).test_files:
-                            test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
-                            test_file.benchmarking_file_path.unlink(missing_ok=True)
-                        for test_file in function_optimizer.test_files.get_by_type(TestType.CONCOLIC_COVERAGE_TEST).test_files:
-                            test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
-
+                    best_optimization = function_optimizer.optimize_function()
                     if is_successful(best_optimization):
                         optimizations_found += 1
                     else:
@@ -190,10 +179,20 @@ def run(self) -> None:
             elif self.args.all:
                 logger.info("✨ All functions have been optimized! ✨")
         finally:
+            if function_optimizer:
+                for test_file in function_optimizer.test_files.get_by_type(TestType.GENERATED_REGRESSION).test_files:
+                    test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
+                    test_file.benchmarking_file_path.unlink(missing_ok=True)
+                for test_file in function_optimizer.test_files.get_by_type(TestType.EXISTING_UNIT_TEST).test_files:
+                    test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
+                    test_file.benchmarking_file_path.unlink(missing_ok=True)
+                for test_file in function_optimizer.test_files.get_by_type(TestType.CONCOLIC_COVERAGE_TEST).test_files:
+                    test_file.instrumented_behavior_file_path.unlink(missing_ok=True)
+                if function_optimizer.test_cfg.concolic_test_root_dir:
+                    shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True)
             if hasattr(get_run_tmp_file, "tmpdir"):
                 get_run_tmp_file.tmpdir.cleanup()
-            if self.test_cfg.concolic_test_root_dir:
-                shutil.rmtree(self.test_cfg.concolic_test_root_dir, ignore_errors=True)
+
 
 
 def run_with_args(args: Namespace) -> None: