diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index c71b9dc29..1bc4d66bb 100644 --- a/codeflash/code_utils/code_utils.py +++ b/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 diff --git a/codeflash/code_utils/config_parser.py b/codeflash/code_utils/config_parser.py index b3f3495d6..d322068cf 100644 --- a/codeflash/code_utils/config_parser.py +++ b/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: diff --git a/tests/code_utils/test_code_utils.py b/tests/code_utils/test_code_utils.py new file mode 100644 index 000000000..73a8e8b2f --- /dev/null +++ b/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()