Test cache revival

codeflash/code_utils/code_utils.py

@@ -238,7 +238,9 @@ def get_all_function_names(code: str) -> tuple[bool, list[str]]:
     return True, function_names
 
 
-def get_run_tmp_file(file_path: Path) -> Path:
+def get_run_tmp_file(file_path: Path | str) -> Path:
+    if isinstance(file_path, str):
+        file_path = Path(file_path)
     if not hasattr(get_run_tmp_file, "tmpdir"):
         get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_")
     return Path(get_run_tmp_file.tmpdir.name) / file_path

codeflash/discovery/discover_unit_tests.py

@@ -54,7 +54,6 @@ class TestsCache:
     def __init__(self) -> None:
         self.connection = sqlite3.connect(codeflash_cache_db)
         self.cur = self.connection.cursor()
-
         self.cur.execute(
             """
             CREATE TABLE IF NOT EXISTS discovered_tests(
@@ -76,7 +75,8 @@ def __init__(self) -> None:
             ON discovered_tests (file_path, file_hash)
             """
         )
-        self._memory_cache = {}
+
+        self.memory_cache = {}
 
     def insert_test(
         self,
@@ -107,25 +107,36 @@ def insert_test(
         )
         self.connection.commit()
 
-    def get_tests_for_file(self, file_path: str, file_hash: str) -> list[FunctionCalledInTest]:
+    def get_function_to_test_map_for_file(
+        self, file_path: str, file_hash: str
+    ) -> dict[str, set[FunctionCalledInTest]] | None:
         cache_key = (file_path, file_hash)
-        if cache_key in self._memory_cache:
-            return self._memory_cache[cache_key]
+        if cache_key in self.memory_cache:
+            return self.memory_cache[cache_key]
+
         self.cur.execute("SELECT * FROM discovered_tests WHERE file_path = ? AND file_hash = ?", (file_path, file_hash))
-        result = [
-            FunctionCalledInTest(
+        rows = self.cur.fetchall()
+        if not rows:
+            return None
+
+        function_to_test_map = defaultdict(set)
+
+        for row in rows:
+            qualified_name_with_modules_from_root = row[2]
+            function_called_in_test = FunctionCalledInTest(
                 tests_in_file=TestsInFile(
                     test_file=Path(row[0]), test_class=row[4], test_function=row[5], test_type=TestType(int(row[6]))
                 ),
                 position=CodePosition(line_no=row[7], col_no=row[8]),
             )
-            for row in self.cur.fetchall()
-        ]
-        self._memory_cache[cache_key] = result
+            function_to_test_map[qualified_name_with_modules_from_root].add(function_called_in_test)
+
+        result = dict(function_to_test_map)
+        self.memory_cache[cache_key] = result
         return result
 
     @staticmethod
-    def compute_file_hash(path: str) -> str:
+    def compute_file_hash(path: str | Path) -> str:
         h = hashlib.sha256(usedforsecurity=False)
         with Path(path).open("rb") as f:
             while True:
@@ -378,7 +389,7 @@ def discover_tests_pytest(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
 
@@ -415,9 +426,11 @@ def discover_tests_pytest(
                 f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}\n {error_section}"
             )
             if "ModuleNotFoundError" in result.stdout:
-                match = ImportErrorPattern.search(result.stdout).group()
-                panel = Panel(Text.from_markup(f"⚠️  {match} ", style="bold red"), expand=False)
-                console.print(panel)
+                match = ImportErrorPattern.search(result.stdout)
+                if match:
+                    error_message = match.group()
+                    panel = Panel(Text.from_markup(f"⚠️  {error_message} ", style="bold red"), expand=False)
+                    console.print(panel)
 
         elif 0 <= exitcode <= 5:
             logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}")
@@ -454,7 +467,7 @@ def discover_tests_unittest(
     cfg: TestConfig,
     discover_only_these_tests: list[str] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
     tests: unittest.TestSuite = loader.discover(str(tests_root))
@@ -510,9 +523,9 @@ def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
 
 
 def discover_parameters_unittest(function_name: str) -> tuple[bool, str, str | None]:
-    function_name = function_name.split("_")
-    if len(function_name) > 1 and function_name[-1].isdigit():
-        return True, "_".join(function_name[:-1]), function_name[-1]
+    function_parts = function_name.split("_")
+    if len(function_parts) > 1 and function_parts[-1].isdigit():
+        return True, "_".join(function_parts[:-1]), function_parts[-1]
 
     return False, function_name, None
 
@@ -521,7 +534,7 @@ def process_test_files(
     file_to_test_map: dict[Path, list[TestsInFile]],
     cfg: TestConfig,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int, int]:
     import jedi
 
     project_root_path = cfg.project_root_path
@@ -536,29 +549,39 @@ def process_test_files(
     num_discovered_replay_tests = 0
     jedi_project = jedi.Project(path=project_root_path)
 
+    tests_cache = TestsCache()
+
     with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
         progress,
         task_id,
     ):
         for test_file, functions in file_to_test_map.items():
+            file_hash = TestsCache.compute_file_hash(test_file)
+
+            cached_function_to_test_map = tests_cache.get_function_to_test_map_for_file(str(test_file), file_hash)
+
+            if cached_function_to_test_map:
+                for qualified_name, test_set in cached_function_to_test_map.items():
+                    function_to_test_map[qualified_name].update(test_set)
+
+                    for function_called_in_test in test_set:
+                        if function_called_in_test.tests_in_file.test_type == TestType.REPLAY_TEST:
+                            num_discovered_replay_tests += 1
+                        num_discovered_tests += 1
+
+                progress.advance(task_id)
+                continue
             try:
                 script = jedi.Script(path=test_file, project=jedi_project)
                 test_functions = set()
 
-                # Single call to get all names with references and definitions
-                all_names = script.get_names(all_scopes=True, references=True, definitions=True)
+                all_names = script.get_names(all_scopes=True, references=True)
+                all_defs = script.get_names(all_scopes=True, definitions=True)
+                all_names_top = script.get_names(all_scopes=True)
 
-                # Filter once and create lookup dictionaries
-                top_level_functions = {}
-                top_level_classes = {}
-                all_defs = []
+                top_level_functions = {name.name: name for name in all_names_top if name.type == "function"}
+                top_level_classes = {name.name: name for name in all_names_top if name.type == "class"}
 
-                for name in all_names:
-                    if name.type == "function":
-                        top_level_functions[name.name] = name
-                        all_defs.append(name)
-                    elif name.type == "class":
-                        top_level_classes[name.name] = name
             except Exception as e:
                 logger.debug(f"Failed to get jedi script for {test_file}: {e}")
                 progress.advance(task_id)
@@ -680,6 +703,18 @@ def process_test_files(
                                     position=CodePosition(line_no=name.line, col_no=name.column),
                                 )
                             )
+                            tests_cache.insert_test(
+                                file_path=str(test_file),
+                                file_hash=file_hash,
+                                qualified_name_with_modules_from_root=qualified_name_with_modules_from_root,
+                                function_name=scope,
+                                test_class=test_func.test_class or "",
+                                test_function=scope_test_function,
+                                test_type=test_func.test_type,
+                                line_number=name.line,
+                                col_number=name.column,
+                            )
+
                             if test_func.test_type == TestType.REPLAY_TEST:
                                 num_discovered_replay_tests += 1
 
@@ -690,4 +725,6 @@ def process_test_files(
 
             progress.advance(task_id)
 
+    tests_cache.close()
+
     return dict(function_to_test_map), num_discovered_tests, num_discovered_replay_tests

codeflash/optimization/optimizer.py

@@ -238,15 +238,14 @@ def discover_tests(
         from codeflash.discovery.discover_unit_tests import discover_unit_tests
 
         console.rule()
-        with progress_bar("Discovering existing function tests..."):
-            start_time = time.time()
-            function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
-                self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
-            )
-            console.rule()
-            logger.info(
-                f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
-            )
+        start_time = time.time()
+        function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
+            self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
+        )
+        console.rule()
+        logger.info(
+            f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
+        )
         console.rule()
         ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
         return function_to_tests, num_discovered_tests