Merge branch 'test_cache_revival' of github.com:codeflash-ai/codeflash into test_cache_revival

Author: Codeflash Bot

codeflash/code_utils/code_extractor.py

@@ -528,16 +528,29 @@ def add_needed_imports_from_module(
 
     try:
         for mod in gatherer.module_imports:
+            # Skip __future__ imports as they cannot be imported directly
+            # __future__ imports should only be imported with specific objects i.e from __future__ import annotations
+            if mod == "__future__":
+                continue
             if mod not in dotted_import_collector.imports:
                 AddImportsVisitor.add_needed_import(dst_context, mod)
             RemoveImportsVisitor.remove_unused_import(dst_context, mod)
+        aliased_objects = set()
+        for mod, alias_pairs in gatherer.alias_mapping.items():
+            for alias_pair in alias_pairs:
+                if alias_pair[0] and alias_pair[1]:  # Both name and alias exist
+                    aliased_objects.add(f"{mod}.{alias_pair[0]}")
+
         for mod, obj_seq in gatherer.object_mapping.items():
             for obj in obj_seq:
                 if (
                     f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
                 ):
                     continue  # Skip adding imports for helper functions already in the context
 
+                if f"{mod}.{obj}" in aliased_objects:
+                    continue
+
                 # Handle star imports by resolving them to actual symbol names
                 if obj == "*":
                     resolved_symbols = resolve_star_import(mod, project_root)
@@ -559,6 +572,8 @@ def add_needed_imports_from_module(
         return dst_module_code
 
     for mod, asname in gatherer.module_aliases.items():
+        if not asname:
+            continue
         if f"{mod}.{asname}" not in dotted_import_collector.imports:
             AddImportsVisitor.add_needed_import(dst_context, mod, asname=asname)
         RemoveImportsVisitor.remove_unused_import(dst_context, mod, asname=asname)
@@ -568,12 +583,16 @@ def add_needed_imports_from_module(
             if f"{mod}.{alias_pair[0]}" in helper_functions_fqn:
                 continue
 
+            if not alias_pair[0] or not alias_pair[1]:
+                continue
+
             if f"{mod}.{alias_pair[1]}" not in dotted_import_collector.imports:
                 AddImportsVisitor.add_needed_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
             RemoveImportsVisitor.remove_unused_import(dst_context, mod, alias_pair[0], asname=alias_pair[1])
 
     try:
-        transformed_module = AddImportsVisitor(dst_context).transform_module(parsed_dst_module)
+        add_imports_visitor = AddImportsVisitor(dst_context)
+        transformed_module = add_imports_visitor.transform_module(parsed_dst_module)
         transformed_module = RemoveImportsVisitor(dst_context).transform_module(transformed_module)
         return transformed_module.code.lstrip("\n")
     except Exception as e:

codeflash/code_utils/code_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import configparser
 import difflib
 import os
 import re
@@ -15,10 +16,12 @@
 import tomlkit
 
 from codeflash.cli_cmds.console import logger, paneled_text
-from codeflash.code_utils.config_parser import find_pyproject_toml
+from codeflash.code_utils.config_parser import find_pyproject_toml, get_all_closest_config_files
 
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
 
+BLACKLIST_ADDOPTS = ("--benchmark", "--sugar", "--codespeed", "--cov", "--profile", "--junitxml", "-n")
+
 
 def unified_diff_strings(code1: str, code2: str, fromfile: str = "original", tofile: str = "modified") -> str:
     """Return the unified diff between two code strings as a single string.
@@ -81,42 +84,105 @@ def create_rank_dictionary_compact(int_array: list[int]) -> dict[int, int]:
     return {original_index: rank for rank, original_index in enumerate(sorted_indices)}
 
 
-@contextmanager
-def custom_addopts() -> None:
-    pyproject_file = find_pyproject_toml()
-    original_content = None
-    non_blacklist_plugin_args = ""
-
+def filter_args(addopts_args: list[str]) -> list[str]:
+    # Convert BLACKLIST_ADDOPTS to a set for faster lookup of simple matches
+    # But keep tuple for startswith
+    blacklist = BLACKLIST_ADDOPTS
+    # Precompute the length for re-use
+    n = len(addopts_args)
+    filtered_args = []
+    i = 0
+    while i < n:
+        current_arg = addopts_args[i]
+        if current_arg.startswith(blacklist):
+            i += 1
+            if i < n and not addopts_args[i].startswith("-"):
+                i += 1
+        else:
+            filtered_args.append(current_arg)
+            i += 1
+    return filtered_args
+
+
+def modify_addopts(config_file: Path) -> tuple[str, bool]:  # noqa : PLR0911
+    file_type = config_file.suffix.lower()
+    filename = config_file.name
+    config = None
+    if file_type not in {".toml", ".ini", ".cfg"} or not config_file.exists():
+        return "", False
+    # Read original file
+    with Path.open(config_file, encoding="utf-8") as f:
+        content = f.read()
     try:
-        # Read original file
-        if pyproject_file.exists():
-            with Path.open(pyproject_file, encoding="utf-8") as f:
-                original_content = f.read()
-                data = tomlkit.parse(original_content)
-            # Backup original addopts
+        if filename == "pyproject.toml":
+            # use tomlkit
+            data = tomlkit.parse(content)
             original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
             # nothing to do if no addopts present
-            if original_addopts != "" and isinstance(original_addopts, list):
-                original_addopts = [x.strip() for x in original_addopts]
-                non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ")
-                non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""]
-                if non_blacklist_plugin_args != original_addopts:
-                    data["tool"]["pytest"]["ini_options"]["addopts"] = non_blacklist_plugin_args
-                    # Write modified file
-                    with Path.open(pyproject_file, "w", encoding="utf-8") as f:
-                        f.write(tomlkit.dumps(data))
+            if original_addopts == "":
+                return content, False
+            if isinstance(original_addopts, list):
+                original_addopts = " ".join(original_addopts)
+            original_addopts = original_addopts.replace("=", " ")
+            addopts_args = (
+                original_addopts.split()
+            )  # any number of space characters as delimiter, doesn't look at = which is fine
+        else:
+            # use configparser
+            config = configparser.ConfigParser()
+            config.read_string(content)
+            data = {section: dict(config[section]) for section in config.sections()}
+            if config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"}:
+                original_addopts = data.get("pytest", {}).get("addopts", "")  # should only be a string
+            else:
+                original_addopts = data.get("tool:pytest", {}).get("addopts", "")  # should only be a string
+            original_addopts = original_addopts.replace("=", " ")
+            addopts_args = original_addopts.split()
+        new_addopts_args = filter_args(addopts_args)
+        if new_addopts_args == addopts_args:
+            return content, False
+        # change addopts now
+        if file_type == ".toml":
+            data["tool"]["pytest"]["ini_options"]["addopts"] = " ".join(new_addopts_args)
+            # Write modified file
+            with Path.open(config_file, "w", encoding="utf-8") as f:
+                f.write(tomlkit.dumps(data))
+                return content, True
+        elif config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"}:
+            config.set("pytest", "addopts", " ".join(new_addopts_args))
+            # Write modified file
+            with Path.open(config_file, "w", encoding="utf-8") as f:
+                config.write(f)
+                return content, True
+        else:
+            config.set("tool:pytest", "addopts", " ".join(new_addopts_args))
+            # Write modified file
+            with Path.open(config_file, "w", encoding="utf-8") as f:
+                config.write(f)
+                return content, True
+
+    except Exception:
+        logger.debug("Trouble parsing")
+        return content, False  # not modified
+
+
+@contextmanager
+def custom_addopts() -> None:
+    closest_config_files = get_all_closest_config_files()
+
+    original_content = {}
 
+    try:
+        for config_file in closest_config_files:
+            original_content[config_file] = modify_addopts(config_file)
         yield
 
     finally:
         # Restore original file
-        if (
-            original_content
-            and pyproject_file.exists()
-            and tuple(original_addopts) not in {(), tuple(non_blacklist_plugin_args)}
-        ):
-            with Path.open(pyproject_file, "w", encoding="utf-8") as f:
-                f.write(original_content)
+        for file, (content, was_modified) in original_content.items():
+            if was_modified:
+                with Path.open(file, "w", encoding="utf-8") as f:
+                    f.write(content)
 
 
 @contextmanager

codeflash/code_utils/config_parser.py

@@ -5,6 +5,9 @@
 
 import tomlkit
 
+PYPROJECT_TOML_CACHE = {}
+ALL_CONFIG_FILES = {}  # map path to closest config file
+
 
 def find_pyproject_toml(config_file: Path | None = None) -> Path:
     # Find the pyproject.toml file on the root of the project
@@ -19,10 +22,15 @@ def find_pyproject_toml(config_file: Path | None = None) -> Path:
             raise ValueError(msg)
         return config_file
     dir_path = Path.cwd()
-
+    cur_path = dir_path
+    # see if it was encountered before in search
+    if cur_path in PYPROJECT_TOML_CACHE:
+        return PYPROJECT_TOML_CACHE[cur_path]
+    # map current path to closest file
     while dir_path != dir_path.parent:
         config_file = dir_path / "pyproject.toml"
         if config_file.exists():
+            PYPROJECT_TOML_CACHE[cur_path] = config_file
             return config_file
         # Search for pyproject.toml in the parent directories
         dir_path = dir_path.parent
@@ -31,6 +39,33 @@ def find_pyproject_toml(config_file: Path | None = None) -> Path:
     raise ValueError(msg)
 
 
+def get_all_closest_config_files() -> list[Path]:
+    all_closest_config_files = []
+    for file_type in ["pyproject.toml", "pytest.ini", ".pytest.ini", "tox.ini", "setup.cfg"]:
+        closest_config_file = find_closest_config_file(file_type)
+        if closest_config_file:
+            all_closest_config_files.append(closest_config_file)
+    return all_closest_config_files
+
+
+def find_closest_config_file(file_type: str) -> Path | None:
+    # Find the closest pyproject.toml, pytest.ini, tox.ini, or setup.cfg file on the root of the project
+    dir_path = Path.cwd()
+    cur_path = dir_path
+    if cur_path in ALL_CONFIG_FILES and file_type in ALL_CONFIG_FILES[cur_path]:
+        return ALL_CONFIG_FILES[cur_path][file_type]
+    while dir_path != dir_path.parent:
+        config_file = dir_path / file_type
+        if config_file.exists():
+            if cur_path not in ALL_CONFIG_FILES:
+                ALL_CONFIG_FILES[cur_path] = {}
+            ALL_CONFIG_FILES[cur_path][file_type] = config_file
+            return config_file
+        # Search for pyproject.toml in the parent directories
+        dir_path = dir_path.parent
+    return None
+
+
 def find_conftest_files(test_paths: list[Path]) -> list[Path]:
     list_of_conftest_files = set()
     for test_path in test_paths:

codeflash/discovery/discover_unit_tests.py

@@ -67,10 +67,35 @@ class TestFunction:
 
 
 class TestsCache:
+    SCHEMA_VERSION = 1  # Increment this when schema changes
+
     def __init__(self, project_root_path: str | Path) -> None:
         self.project_root_path = Path(project_root_path).resolve().as_posix()
         self.connection = sqlite3.connect(codeflash_cache_db)
         self.cur = self.connection.cursor()
+
+        self.cur.execute(
+            """
+            CREATE TABLE IF NOT EXISTS schema_version(
+                version INTEGER PRIMARY KEY
+            )
+            """
+        )
+
+        self.cur.execute("SELECT version FROM schema_version")
+        result = self.cur.fetchone()
+        current_version = result[0] if result else None
+
+        if current_version != self.SCHEMA_VERSION:
+            logger.debug(
+                f"Schema version mismatch (current: {current_version}, expected: {self.SCHEMA_VERSION}). Recreating tables."
+            )
+            self.cur.execute("DROP TABLE IF EXISTS discovered_tests")
+            self.cur.execute("DROP INDEX IF EXISTS idx_discovered_tests_project_file_path_hash")
+            self.cur.execute("DELETE FROM schema_version")
+            self.cur.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,))
+            self.connection.commit()
+
         self.cur.execute(
             """
             CREATE TABLE IF NOT EXISTS discovered_tests(
@@ -158,14 +183,16 @@ def get_function_to_test_map_for_file(
         return result
 
     @staticmethod
-    def compute_file_hash(path: str | Path) -> str:
+    def compute_file_hash(path: Path) -> str:
         h = hashlib.sha256(usedforsecurity=False)
-        with Path(path).open("rb") as f:
+        with path.open("rb", buffering=0) as f:
+            buf = bytearray(8192)
+            mv = memoryview(buf)
             while True:
-                chunk = f.read(8192)
-                if not chunk:
+                n = f.readinto(mv)
+                if n == 0:
                     break
-                h.update(chunk)
+                h.update(mv[:n])
         return h.hexdigest()
 
     def close(self) -> None:
@@ -488,13 +515,13 @@ def discover_tests_pytest(
 
 def discover_tests_unittest(
     cfg: TestConfig,
-    discover_only_these_tests: list[str] | None = None,
+    discover_only_these_tests: list[Path] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
 ) -> 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))
-    file_to_test_map: defaultdict[str, list[TestsInFile]] = defaultdict(list)
+    file_to_test_map: defaultdict[Path, list[TestsInFile]] = defaultdict(list)
 
     def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
         _test_function, _test_module, _test_suite_name = (
@@ -506,7 +533,7 @@ def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
         _test_module_path = Path(_test_module.replace(".", os.sep)).with_suffix(".py")
         _test_module_path = tests_root / _test_module_path
         if not _test_module_path.exists() or (
-            discover_only_these_tests and str(_test_module_path) not in discover_only_these_tests
+            discover_only_these_tests and _test_module_path not in discover_only_these_tests
         ):
             return None
         if "__replay_test" in str(_test_module_path):
@@ -516,10 +543,7 @@ def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
         else:
             test_type = TestType.EXISTING_UNIT_TEST
         return TestsInFile(
-            test_file=str(_test_module_path),
-            test_function=_test_function,
-            test_type=test_type,
-            test_class=_test_suite_name,
+            test_file=_test_module_path, test_function=_test_function, test_type=test_type, test_class=_test_suite_name
         )
 
     for _test_suite in tests._tests:
@@ -537,11 +561,11 @@ def get_test_details(_test: unittest.TestCase) -> TestsInFile | None:
                             continue
                         details = get_test_details(test_2)
                         if details is not None:
-                            file_to_test_map[str(details.test_file)].append(details)
+                            file_to_test_map[details.test_file].append(details)
                 else:
                     details = get_test_details(test)
                     if details is not None:
-                        file_to_test_map[str(details.test_file)].append(details)
+                        file_to_test_map[details.test_file].append(details)
     return process_test_files(file_to_test_map, cfg, functions_to_optimize)