|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import ast |
| 4 | +import configparser |
4 | 5 | import difflib |
5 | 6 | import os |
6 | 7 | import re |
|
15 | 16 | import tomlkit |
16 | 17 |
|
17 | 18 | from codeflash.cli_cmds.console import logger, paneled_text |
18 | | -from codeflash.code_utils.config_parser import find_pyproject_toml |
| 19 | +from codeflash.code_utils.config_parser import find_pyproject_toml, get_all_closest_config_files |
19 | 20 |
|
20 | 21 | ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE) |
21 | 22 |
|
| 23 | +BLACKLIST_ADDOPTS = ("--benchmark", "--sugar", "--codespeed", "--cov", "--profile", "--junitxml", "-n") |
| 24 | + |
22 | 25 |
|
23 | 26 | def unified_diff_strings(code1: str, code2: str, fromfile: str = "original", tofile: str = "modified") -> str: |
24 | 27 | """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]: |
81 | 84 | return {original_index: rank for rank, original_index in enumerate(sorted_indices)} |
82 | 85 |
|
83 | 86 |
|
84 | | -@contextmanager |
85 | | -def custom_addopts() -> None: |
86 | | - pyproject_file = find_pyproject_toml() |
87 | | - original_content = None |
88 | | - non_blacklist_plugin_args = "" |
89 | | - |
| 87 | +def filter_args(addopts_args: list[str]) -> list[str]: |
| 88 | + # Convert BLACKLIST_ADDOPTS to a set for faster lookup of simple matches |
| 89 | + # But keep tuple for startswith |
| 90 | + blacklist = BLACKLIST_ADDOPTS |
| 91 | + # Precompute the length for re-use |
| 92 | + n = len(addopts_args) |
| 93 | + filtered_args = [] |
| 94 | + i = 0 |
| 95 | + while i < n: |
| 96 | + current_arg = addopts_args[i] |
| 97 | + if current_arg.startswith(blacklist): |
| 98 | + i += 1 |
| 99 | + if i < n and not addopts_args[i].startswith("-"): |
| 100 | + i += 1 |
| 101 | + else: |
| 102 | + filtered_args.append(current_arg) |
| 103 | + i += 1 |
| 104 | + return filtered_args |
| 105 | + |
| 106 | + |
| 107 | +def modify_addopts(config_file: Path) -> tuple[str, bool]: # noqa : PLR0911 |
| 108 | + file_type = config_file.suffix.lower() |
| 109 | + filename = config_file.name |
| 110 | + config = None |
| 111 | + if file_type not in {".toml", ".ini", ".cfg"} or not config_file.exists(): |
| 112 | + return "", False |
| 113 | + # Read original file |
| 114 | + with Path.open(config_file, encoding="utf-8") as f: |
| 115 | + content = f.read() |
90 | 116 | try: |
91 | | - # Read original file |
92 | | - if pyproject_file.exists(): |
93 | | - with Path.open(pyproject_file, encoding="utf-8") as f: |
94 | | - original_content = f.read() |
95 | | - data = tomlkit.parse(original_content) |
96 | | - # Backup original addopts |
| 117 | + if filename == "pyproject.toml": |
| 118 | + # use tomlkit |
| 119 | + data = tomlkit.parse(content) |
97 | 120 | original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "") |
98 | 121 | # nothing to do if no addopts present |
99 | | - if original_addopts != "" and isinstance(original_addopts, list): |
100 | | - original_addopts = [x.strip() for x in original_addopts] |
101 | | - non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ") |
102 | | - non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""] |
103 | | - if non_blacklist_plugin_args != original_addopts: |
104 | | - data["tool"]["pytest"]["ini_options"]["addopts"] = non_blacklist_plugin_args |
105 | | - # Write modified file |
106 | | - with Path.open(pyproject_file, "w", encoding="utf-8") as f: |
107 | | - f.write(tomlkit.dumps(data)) |
| 122 | + if original_addopts == "": |
| 123 | + return content, False |
| 124 | + if isinstance(original_addopts, list): |
| 125 | + original_addopts = " ".join(original_addopts) |
| 126 | + original_addopts = original_addopts.replace("=", " ") |
| 127 | + addopts_args = ( |
| 128 | + original_addopts.split() |
| 129 | + ) # any number of space characters as delimiter, doesn't look at = which is fine |
| 130 | + else: |
| 131 | + # use configparser |
| 132 | + config = configparser.ConfigParser() |
| 133 | + config.read_string(content) |
| 134 | + data = {section: dict(config[section]) for section in config.sections()} |
| 135 | + if config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"}: |
| 136 | + original_addopts = data.get("pytest", {}).get("addopts", "") # should only be a string |
| 137 | + else: |
| 138 | + original_addopts = data.get("tool:pytest", {}).get("addopts", "") # should only be a string |
| 139 | + original_addopts = original_addopts.replace("=", " ") |
| 140 | + addopts_args = original_addopts.split() |
| 141 | + new_addopts_args = filter_args(addopts_args) |
| 142 | + if new_addopts_args == addopts_args: |
| 143 | + return content, False |
| 144 | + # change addopts now |
| 145 | + if file_type == ".toml": |
| 146 | + data["tool"]["pytest"]["ini_options"]["addopts"] = " ".join(new_addopts_args) |
| 147 | + # Write modified file |
| 148 | + with Path.open(config_file, "w", encoding="utf-8") as f: |
| 149 | + f.write(tomlkit.dumps(data)) |
| 150 | + return content, True |
| 151 | + elif config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"}: |
| 152 | + config.set("pytest", "addopts", " ".join(new_addopts_args)) |
| 153 | + # Write modified file |
| 154 | + with Path.open(config_file, "w", encoding="utf-8") as f: |
| 155 | + config.write(f) |
| 156 | + return content, True |
| 157 | + else: |
| 158 | + config.set("tool:pytest", "addopts", " ".join(new_addopts_args)) |
| 159 | + # Write modified file |
| 160 | + with Path.open(config_file, "w", encoding="utf-8") as f: |
| 161 | + config.write(f) |
| 162 | + return content, True |
| 163 | + |
| 164 | + except Exception: |
| 165 | + logger.debug("Trouble parsing") |
| 166 | + return content, False # not modified |
| 167 | + |
| 168 | + |
| 169 | +@contextmanager |
| 170 | +def custom_addopts() -> None: |
| 171 | + closest_config_files = get_all_closest_config_files() |
| 172 | + |
| 173 | + original_content = {} |
108 | 174 |
|
| 175 | + try: |
| 176 | + for config_file in closest_config_files: |
| 177 | + original_content[config_file] = modify_addopts(config_file) |
109 | 178 | yield |
110 | 179 |
|
111 | 180 | finally: |
112 | 181 | # Restore original file |
113 | | - if ( |
114 | | - original_content |
115 | | - and pyproject_file.exists() |
116 | | - and tuple(original_addopts) not in {(), tuple(non_blacklist_plugin_args)} |
117 | | - ): |
118 | | - with Path.open(pyproject_file, "w", encoding="utf-8") as f: |
119 | | - f.write(original_content) |
| 182 | + for file, (content, was_modified) in original_content.items(): |
| 183 | + if was_modified: |
| 184 | + with Path.open(file, "w", encoding="utf-8") as f: |
| 185 | + f.write(content) |
120 | 186 |
|
121 | 187 |
|
122 | 188 | @contextmanager |
|
0 commit comments