Paths cleanup

codeflash/code_utils/code_utils.py

@@ -6,6 +6,7 @@
 from functools import lru_cache
 from pathlib import Path
 from tempfile import TemporaryDirectory
+import shutil
 
 from codeflash.cli_cmds.console import logger
 
@@ -118,4 +119,7 @@ def has_any_async_functions(code: str) -> bool:
 
 def cleanup_paths(paths: list[Path]) -> None:
     for path in paths:
-        path.unlink(missing_ok=True)
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)

codeflash/optimization/function_optimizer.py

@@ -319,20 +319,8 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                             original_helper_code,
                             self.function_to_optimize.file_path,
                         )
-        for generated_test_path in generated_test_paths:
-            generated_test_path.unlink(missing_ok=True)
-        for generated_perf_test_path in generated_perf_test_paths:
-            generated_perf_test_path.unlink(missing_ok=True)
-        for test_paths in instrumented_unittests_created_for_function:
-            test_paths.unlink(missing_ok=True)
-        for fn in function_to_concolic_tests:
-            for test in function_to_concolic_tests[fn]:
-                if not test.tests_in_file.test_file.parent.exists():
-                    logger.warning(
-                        f"Concolic test directory {test.tests_in_file.test_file.parent} does not exist so could not be deleted."
-                    )
-                shutil.rmtree(test.tests_in_file.test_file.parent, ignore_errors=True)
-                break  # need to delete only one test directory
+        cleanup_paths(generated_test_paths + generated_perf_test_paths)
+
 
         if not best_optimization:
             return Failure(f"No best optimizations found for function {self.function_to_optimize.qualified_name}")

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()