Merge pull request #812 from codeflash-ai/improve-addopts

Improve addopts masking (fixes repos with addopts which interfere with cf)

Author: aseembits93

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,8 @@
 
 import tomlkit
 
+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
@@ -31,6 +33,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:

tests/code_utils/test_code_utils.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import configparser
+import os
+import stat
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import tomlkit
+
+from codeflash.code_utils.code_utils import custom_addopts
+
+def test_custom_addopts_modifies_and_restores_dotini_file(tmp_path: Path) -> None:
+    """Verify that custom_addopts correctly modifies and then restores a pytest.ini file."""
+    # Create a dummy pytest.ini file
+    config_file = tmp_path / ".pytest.ini"
+    original_content = "[pytest]\naddopts = -v --cov=./src -n auto\n"
+    config_file.write_text(original_content)
+
+    # Use patch to mock get_all_closest_config_files
+    os.chdir(tmp_path)
+    with custom_addopts():
+        # Check that the file is modified inside the context
+        modified_content = config_file.read_text()
+        config = configparser.ConfigParser()
+        config.read_string(modified_content)
+        modified_addopts = config.get("pytest", "addopts", fallback="")
+        assert modified_addopts == "-v"
+
+    # Check that the file is restored after exiting the context
+    restored_content = config_file.read_text()
+    assert restored_content.strip() == original_content.strip()
+
+def test_custom_addopts_modifies_and_restores_ini_file(tmp_path: Path) -> None:
+    """Verify that custom_addopts correctly modifies and then restores a pytest.ini file."""
+    # Create a dummy pytest.ini file
+    config_file = tmp_path / "pytest.ini"
+    original_content = "[pytest]\naddopts = -v --cov=./src -n auto\n"
+    config_file.write_text(original_content)
+
+    # Use patch to mock get_all_closest_config_files
+    os.chdir(tmp_path)
+    with custom_addopts():
+        # Check that the file is modified inside the context
+        modified_content = config_file.read_text()
+        config = configparser.ConfigParser()
+        config.read_string(modified_content)
+        modified_addopts = config.get("pytest", "addopts", fallback="")
+        assert modified_addopts == "-v"
+
+    # Check that the file is restored after exiting the context
+    restored_content = config_file.read_text()
+    assert restored_content.strip() == original_content.strip()
+
+
+def test_custom_addopts_modifies_and_restores_toml_file(tmp_path: Path) -> None:
+    """Verify that custom_addopts correctly modifies and then restores a pyproject.toml file."""
+    # Create a dummy pyproject.toml file
+    config_file = tmp_path / "pyproject.toml"
+    os.chdir(tmp_path)
+    original_addopts = "-v --cov=./src --junitxml=report.xml"
+    original_content_dict = {
+        "tool": {"pytest": {"ini_options": {"addopts": original_addopts}}}
+    }
+    original_content = tomlkit.dumps(original_content_dict)
+    config_file.write_text(original_content)
+
+    # Use patch to mock get_all_closest_config_files
+    os.chdir(tmp_path)
+    with custom_addopts():
+        # Check that the file is modified inside the context
+        modified_content = config_file.read_text()
+        modified_data = tomlkit.parse(modified_content)
+        modified_addopts = modified_data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
+        assert modified_addopts == "-v"
+
+    # Check that the file is restored after exiting the context
+    restored_content = config_file.read_text()
+    assert restored_content.strip() == original_content.strip()
+
+
+def test_custom_addopts_handles_no_addopts(tmp_path: Path) -> None:
+    """Ensure custom_addopts doesn't fail when a config file has no addopts."""
+    # Create a dummy pytest.ini file without addopts
+    config_file = tmp_path / "pytest.ini"
+    original_content = "[pytest]\n# no addopts here\n"
+    config_file.write_text(original_content)
+
+    os.chdir(tmp_path)
+    with custom_addopts():
+        # The file should not be modified
+        content_inside_context = config_file.read_text()
+        assert content_inside_context == original_content
+
+    # The file should remain unchanged
+    content_after_context = config_file.read_text()
+    assert content_after_context == original_content
+
+def test_custom_addopts_handles_no_relevant_files(tmp_path: Path) -> None:
+    """Ensure custom_addopts runs without error when no config files are found."""
+    # No config files created in tmp_path
+
+    os.chdir(tmp_path)
+    # This should execute without raising any exceptions
+    with custom_addopts():
+        pass
+    # No assertions needed, the test passes if no exceptions were raised
+
+
+def test_custom_addopts_toml_without_pytest_section(tmp_path: Path) -> None:
+    """Verify custom_addopts doesn't fail with a toml file missing a [tool.pytest] section."""
+    config_file = tmp_path / "pyproject.toml"
+    original_content_dict = {"tool": {"other_tool": {"key": "value"}}}
+    original_content = tomlkit.dumps(original_content_dict)
+    config_file.write_text(original_content)
+
+    os.chdir(tmp_path)
+    with custom_addopts():
+        content_inside_context = config_file.read_text()
+        assert content_inside_context == original_content
+
+    content_after_context = config_file.read_text()
+    assert content_after_context == original_content
+
+
+def test_custom_addopts_ini_without_pytest_section(tmp_path: Path) -> None:
+    """Verify custom_addopts doesn't fail with an ini file missing a [pytest] section."""
+    config_file = tmp_path / "pytest.ini"
+    original_content = "[other_section]\nkey = value\n"
+    config_file.write_text(original_content)
+
+    os.chdir(tmp_path)
+    with custom_addopts():
+        content_inside_context = config_file.read_text()
+        assert content_inside_context == original_content
+
+    content_after_context = config_file.read_text()
+    assert content_after_context == original_content
+
+
+def test_custom_addopts_with_multiple_config_files(tmp_path: Path) -> None:
+    """Verify custom_addopts modifies and restores all found config files."""
+    os.chdir(tmp_path)
+
+    # Create pytest.ini
+    ini_file = tmp_path / "pytest.ini"
+    ini_original_content = "[pytest]\naddopts = -v --cov\n"
+    ini_file.write_text(ini_original_content)
+
+    # Create pyproject.toml
+    toml_file = tmp_path / "pyproject.toml"
+    toml_original_addopts = "-s -n auto"
+    toml_original_content_dict = {
+        "tool": {"pytest": {"ini_options": {"addopts": toml_original_addopts}}}
+    }
+    toml_original_content = tomlkit.dumps(toml_original_content_dict)
+    toml_file.write_text(toml_original_content)
+
+    with custom_addopts():
+        # Check INI file modification
+        ini_modified_content = ini_file.read_text()
+        config = configparser.ConfigParser()
+        config.read_string(ini_modified_content)
+        assert config.get("pytest", "addopts", fallback="") == "-v"
+
+        # Check TOML file modification
+        toml_modified_content = toml_file.read_text()
+        modified_data = tomlkit.parse(toml_modified_content)
+        modified_addopts = modified_data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
+        assert modified_addopts == "-s"
+
+    # Check that both files are restored
+    assert ini_file.read_text().strip() == ini_original_content.strip()
+    assert toml_file.read_text().strip() == toml_original_content.strip()
+
+
+def test_custom_addopts_restores_on_exception(tmp_path: Path) -> None:
+    """Ensure config file is restored even if an exception occurs inside the context."""
+    config_file = tmp_path / "pytest.ini"
+    original_content = "[pytest]\naddopts = -v --cov\n"
+    config_file.write_text(original_content)
+
+    os.chdir(tmp_path)
+    with pytest.raises(ValueError, match="Test exception"):
+        with custom_addopts():
+            raise ValueError("Test exception")
+
+    restored_content = config_file.read_text()
+    assert restored_content.strip() == original_content.strip()