cleanup paths in optimizer.py

Author: DemetrioReyes

codeflash/optimization/function_optimizer.py

@@ -319,11 +319,7 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                             original_helper_code,
                             self.function_to_optimize.file_path,
                         )
-                        
-        # lista filepath
-        filespaths = generated_test_paths + generated_perf_test_paths
-        # remove all test files
-        cleanup_paths(filespaths)
+        cleanup_paths(generated_test_paths + generated_perf_test_paths)
 
 
         if not best_optimization:

codeflash/optimization/optimizer.py

@@ -20,6 +20,7 @@
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.telemetry.posthog_cf import ph
 from codeflash.verification.verification_utils import TestConfig
+from codeflash.code_utils.code_utils import cleanup_paths
 
 if TYPE_CHECKING:
     from argparse import Namespace
@@ -177,16 +178,22 @@ def run(self) -> None:
                 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)
+                list_of_paths_to_delete = [
+                    test_file.instrumented_behavior_file_path
+                    for test_type in [
+                        TestType.GENERATED_REGRESSION,
+                        TestType.EXISTING_UNIT_TEST,
+                        TestType.CONCOLIC_COVERAGE_TEST,
+                    ]
+                    for test_file in function_optimizer.test_files.get_by_type(test_type).test_files
+                ] + [
+                    test_file.benchmarking_file_path
+                    for test_type in [TestType.GENERATED_REGRESSION, TestType.EXISTING_UNIT_TEST]
+                    for test_file in function_optimizer.test_files.get_by_type(test_type).test_files
+                ] + [function_optimizer.test_cfg.concolic_test_root_dir]
+
+                cleanup_paths(list_of_paths_to_delete)
+
             if hasattr(get_run_tmp_file, "tmpdir"):
                 get_run_tmp_file.tmpdir.cleanup()