consolidate logic around path cleanup

Author: KRRT7

codeflash/code_utils/code_utils.py

@@ -2,6 +2,7 @@
 
 import ast
 import os
+import shutil
 import site
 from functools import lru_cache
 from pathlib import Path
@@ -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

@@ -345,20 +345,6 @@ 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
 
         if not best_optimization:
             return Failure(f"No best optimizations found for function {self.function_to_optimize.qualified_name}")
@@ -1242,3 +1228,25 @@ def generate_and_instrument_tests(
                 zip(generated_test_paths, generated_perf_test_paths)
             )
         ]
+
+    def cleanup_generated_files(self) -> None:
+        paths_to_cleanup = (
+            [
+                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 self.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 self.test_files.get_by_type(test_type).test_files
+            ]
+            + [self.test_cfg.concolic_test_root_dir]
+        )
+        cleanup_paths(paths_to_cleanup)
+        if hasattr(get_run_tmp_file, "tmpdir"):
+            get_run_tmp_file.tmpdir.cleanup()

codeflash/optimization/optimizer.py

@@ -2,7 +2,6 @@
 
 import ast
 import os
-import shutil
 import tempfile
 import time
 from collections import defaultdict
@@ -19,7 +18,7 @@
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.checkpoint import CodeflashRunCheckpoint, ask_should_use_checkpoint_get_functions
 from codeflash.code_utils.code_replacer import normalize_code, normalize_node
-from codeflash.code_utils.code_utils import get_run_tmp_file
+from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file
 from codeflash.code_utils.static_analysis import analyze_imported_modules, get_first_top_level_function_or_method_ast
 from codeflash.discovery.discover_unit_tests import discover_unit_tests
 from codeflash.discovery.functions_to_optimize import get_functions_to_optimize
@@ -266,22 +265,7 @@ 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)
-                if self.args.benchmark:
-                    if self.replay_tests_dir.exists():
-                        shutil.rmtree(self.replay_tests_dir, ignore_errors=True)
-                    trace_file.unlink(missing_ok=True)
-            if hasattr(get_run_tmp_file, "tmpdir"):
-                get_run_tmp_file.tmpdir.cleanup()
+                function_optimizer.cleanup_generated_files()
 
 
 def run_with_args(args: Namespace) -> None: