{{ message }}
-
Notifications
You must be signed in to change notification settings - Fork 26
Expand file tree
/
Copy pathcode_utils.py
More file actions
543 lines (415 loc) · 18.6 KB
/
Copy pathcode_utils.py
File metadata and controls
543 lines (415 loc) · 18.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
from __future__ import annotations
import ast
import configparser
import difflib
import os
import re
import shutil
import site
import sys
from contextlib import contextmanager
from functools import lru_cache
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Generator
import tomlkit
from codeflash.cli_cmds.console import logger, paneled_text
from codeflash.code_utils.config_parser import find_pyproject_toml, get_all_closest_config_files
from codeflash.lsp.helpers import is_LSP_enabled, is_subagent_mode
_INVALID_CHARS_NT = {"<", ">", ":", '"', "|", "?", "*"}
_INVALID_CHARS_UNIX = {"\0"}
ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
BLACKLIST_ADDOPTS = ("--benchmark", "--sugar", "--codespeed", "--cov", "--profile", "--junitxml", "-n")
# Characters that indicate a glob pattern
GLOB_PATTERN_CHARS = frozenset("*?[")
def is_glob_pattern(path_str: str) -> bool:
"""Check if a path string contains glob pattern characters."""
return any(char in path_str for char in GLOB_PATTERN_CHARS)
def normalize_ignore_paths(paths: list[str], base_path: Path | None = None) -> list[Path]:
if base_path is None:
base_path = Path.cwd()
base_path = base_path.resolve()
normalized: set[Path] = set()
for path_str in paths:
if not path_str:
continue
path_str = str(path_str)
if is_glob_pattern(path_str):
# pathlib requires relative glob patterns
path_str = path_str.removeprefix("./")
if path_str.startswith("/"):
path_str = path_str.lstrip("/")
for matched_path in base_path.glob(path_str):
normalized.add(matched_path.resolve())
else:
path_obj = Path(path_str)
if not path_obj.is_absolute():
path_obj = base_path / path_obj
if path_obj.exists():
normalized.add(path_obj.resolve())
return list(normalized)
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.
:param code1: First code string (original).
:param code2: Second code string (modified).
:param fromfile: Label for the first code string.
:param tofile: Label for the second code string.
:return: Unified diff as a string.
"""
code1_lines = code1.splitlines(keepends=True)
code2_lines = code2.splitlines(keepends=True)
diff = difflib.unified_diff(code1_lines, code2_lines, fromfile=fromfile, tofile=tofile, lineterm="")
return "".join(diff)
def choose_weights(**importance: float) -> list[float]:
"""Choose normalized weights from relative importance values.
Example:
choose_weights(runtime=3, diff=1)
-> [0.75, 0.25]
Args:
**importance: keyword args of metric=importance (relative numbers).
Returns:
A list of weights in the same order as the arguments.
"""
total = sum(importance.values())
if total == 0:
raise ValueError("At least one importance value must be > 0")
return [v / total for v in importance.values()]
def normalize_by_max(values: list[float]) -> list[float]:
mx = max(values)
if mx == 0:
return [0.0] * len(values)
return [v / mx for v in values]
def create_score_dictionary_from_metrics(weights: list[float], *metrics: list[float]) -> dict[int, float]:
"""Combine multiple metrics into a single weighted score dictionary.
Each metric is a list of values (smaller = better).
The total score for each index is the weighted sum of its values
across all metrics:
score[index] = Σ (value * weight)
Args:
weights: A list of weights, one per metric. Larger weight = more influence.
*metrics: Lists of values (one list per metric, aligned by index).
Returns:
A dictionary mapping each index to its combined weighted score.
"""
if len(weights) != len(metrics):
raise ValueError("Number of weights must match number of metrics")
combined: dict[int, float] = {}
for weight, metric in zip(weights, metrics):
for idx, value in enumerate(metric):
combined[idx] = combined.get(idx, 0) + value * weight
return combined
def diff_length(a: str, b: str) -> int:
"""Compute the length (in characters) of the unified diff between two strings.
Args:
a (str): Original string.
b (str): Modified string.
Returns:
int: Total number of characters in the diff.
"""
# Split input strings into lines for line-by-line diff
a_lines = a.splitlines(keepends=True)
b_lines = b.splitlines(keepends=True)
# Compute unified diff
diff_lines = list(difflib.unified_diff(a_lines, b_lines, lineterm=""))
# Join all lines with newline to calculate total diff length
diff_text = "\n".join(diff_lines)
return len(diff_text)
def create_rank_dictionary_compact(int_array: list[int]) -> dict[int, int]:
"""Create a dictionary from a list of ints, mapping the original index to its rank.
This version uses a more compact, "Pythonic" implementation.
Args:
int_array: A list of integers.
Returns:
A dictionary where keys are original indices and values are the
rank of the element in ascending order.
"""
# Sort the indices of the array based on their corresponding values
sorted_indices = sorted(range(len(int_array)), key=lambda i: int_array[i])
# Create a dictionary mapping the original index to its rank (its position in the sorted list)
return {original_index: rank for rank, original_index in enumerate(sorted_indices)}
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]:
file_type = config_file.suffix.lower()
filename = config_file.name
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:
if filename == "pyproject.toml":
data = tomlkit.parse(content)
original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
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()
new_addopts_args = filter_args(addopts_args)
if new_addopts_args == addopts_args:
return content, False
data["tool"]["pytest"]["ini_options"]["addopts"] = " ".join(new_addopts_args) # type: ignore[index]
with Path.open(config_file, "w", encoding="utf-8") as f:
f.write(tomlkit.dumps(data))
return content, True
config = configparser.ConfigParser()
config.read_string(content)
ini_data = {section: dict(config[section]) for section in config.sections()}
if config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"}:
original_addopts = ini_data.get("pytest", {}).get("addopts", "")
else:
original_addopts = ini_data.get("tool:pytest", {}).get("addopts", "")
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
section = "pytest" if config_file.name in {"pytest.ini", ".pytest.ini", "tox.ini"} else "tool:pytest"
config.set(section, "addopts", " ".join(new_addopts_args))
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
@contextmanager
def custom_addopts() -> Generator[None, None, 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
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
def add_addopts_to_pyproject() -> Generator[None, None, None]:
pyproject_file = find_pyproject_toml()
original_content: str | None = None
try:
if pyproject_file.exists():
with Path.open(pyproject_file, encoding="utf-8") as f:
original_content = f.read()
data = tomlkit.parse(original_content)
data["tool"]["pytest"] = {} # type: ignore[index]
data["tool"]["pytest"]["ini_options"] = {} # type: ignore[index]
data["tool"]["pytest"]["ini_options"]["addopts"] = [ # type: ignore[index]
"-n=auto",
"-n",
"1",
"-n 1",
"-n 1",
"-n auto",
]
with Path.open(pyproject_file, "w", encoding="utf-8") as f:
f.write(tomlkit.dumps(data))
yield
finally:
if original_content is not None:
with Path.open(pyproject_file, "w", encoding="utf-8") as f:
f.write(original_content)
def encoded_tokens_len(s: str) -> int:
"""Return the approximate length of the encoded tokens.
It's an approximation of BPE encoding (https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf).
"""
return int(len(s) * 0.25)
def get_qualified_name(module_name: str, full_qualified_name: str) -> str:
if not full_qualified_name:
msg = "full_qualified_name cannot be empty"
raise ValueError(msg)
if not full_qualified_name.startswith(module_name):
msg = f"{full_qualified_name} does not start with {module_name}"
raise ValueError(msg)
if module_name == full_qualified_name:
msg = f"{full_qualified_name} is the same as {module_name}"
raise ValueError(msg)
return full_qualified_name[len(module_name) + 1 :]
_PARAMETERIZED_INDEX_RE = re.compile(r"\[(\d+)")
def extract_parameterized_test_index(test_name: str) -> int:
"""Extract the numeric index from a parameterized test name.
Handles formats like ``test[ 0 ]``, ``test[1]``, and
``test[1] input=foo, expected=bar``. Returns 1 when no numeric index is found.
"""
m = _PARAMETERIZED_INDEX_RE.search(test_name)
return int(m.group(1)) if m else 1
def module_name_from_file_path(file_path: Path, project_root_path: Path, *, traverse_up: bool = False) -> str:
try:
relative_path = file_path.resolve().relative_to(project_root_path.resolve())
return relative_path.with_suffix("").as_posix().replace("/", ".")
except ValueError:
if traverse_up:
parent = file_path.parent
while parent not in (project_root_path, parent.parent):
try:
relative_path = file_path.resolve().relative_to(parent.resolve())
return relative_path.with_suffix("").as_posix().replace("/", ".")
except ValueError:
parent = parent.parent
msg = f"File {file_path} is not within the project root {project_root_path}."
raise ValueError(msg) # noqa: B904
def file_path_from_module_name(module_name: str, project_root_path: Path) -> Path:
"""Get file path from module path."""
return project_root_path / (module_name.replace(".", os.sep) + ".py")
@lru_cache(maxsize=100)
def file_name_from_test_module_name(test_module_name: str, base_dir: Path) -> Path | None:
partial_test_class = test_module_name
while partial_test_class:
test_path = file_path_from_module_name(partial_test_class, base_dir)
if (base_dir / test_path).exists():
return base_dir / test_path
partial_test_class = ".".join(partial_test_class.split(".")[:-1])
return None
def get_imports_from_file(
file_path: Path | None = None, file_string: str | None = None, file_ast: ast.AST | None = None
) -> list[ast.Import | ast.ImportFrom]:
assert sum([file_path is not None, file_string is not None, file_ast is not None]) == 1, (
"Must provide exactly one of file_path, file_string, or file_ast"
)
if file_path:
with file_path.open(encoding="utf8") as file:
file_string = file.read()
if file_ast is None:
if file_string is None:
logger.error("file_string cannot be None when file_ast is not provided")
return []
try:
file_ast = ast.parse(file_string)
except SyntaxError as e:
logger.exception(f"Syntax error in code: {e}")
return []
return [node for node in ast.walk(file_ast) if isinstance(node, (ast.Import, ast.ImportFrom))]
def get_all_function_names(code: str) -> tuple[bool, list[str]]:
try:
module = ast.parse(code)
except SyntaxError as e:
logger.exception(f"Syntax error in code: {e}")
return False, []
function_names = [
node.name for node in ast.walk(module) if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
]
return True, function_names
_run_tmpdir: TemporaryDirectory[str] | None = None
_run_tmpdir_path: Path | None = None
def get_run_tmp_file(file_path: Path | str) -> Path:
global _run_tmpdir, _run_tmpdir_path
if isinstance(file_path, str):
file_path = Path(file_path)
if _run_tmpdir_path is None:
_run_tmpdir = TemporaryDirectory(prefix="codeflash_")
_run_tmpdir_path = Path(_run_tmpdir.name).resolve()
return _run_tmpdir_path / file_path
def path_belongs_to_site_packages(file_path: Path) -> bool:
file_path_resolved = file_path.resolve()
site_packages = [Path(p).resolve() for p in site.getsitepackages()]
return any(file_path_resolved.is_relative_to(site_package_path) for site_package_path in site_packages)
def is_class_defined_in_file(class_name: str, file_path: Path) -> bool:
if not file_path.exists():
return False
with file_path.open(encoding="utf8") as file:
source = file.read()
tree = ast.parse(source)
return any(isinstance(node, ast.ClassDef) and node.name == class_name for node in ast.walk(tree))
def validate_python_code(code: str) -> str:
"""Validate a string of Python code by attempting to compile it."""
try:
compile(code, "<string>", "exec")
except SyntaxError as e:
msg = f"Invalid Python code: {e.msg} (line {e.lineno}, column {e.offset})"
raise ValueError(msg) from e
return code
def cleanup_paths(paths: list[Path]) -> None:
for path in paths:
if path and path.exists():
if path.is_dir():
shutil.rmtree(path, ignore_errors=True)
else:
path.unlink(missing_ok=True)
def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
for path, file_content in path_to_content_map.items():
path.write_text(file_content, encoding="utf8")
def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
"""Don't Call it inside the lsp process, it will terminate the lsp server."""
if is_LSP_enabled():
logger.error(message)
return
if is_subagent_mode():
from xml.sax.saxutils import escape
sys.stdout.write(f"<codeflash-error>{escape(message)}</codeflash-error>\n")
sys.exit(1 if error_on_exit else 0)
paneled_text(message, panel_args={"style": "red"})
sys.exit(1 if error_on_exit else 0)
def shorten_pytest_error(pytest_error_string: str) -> str:
return "\n".join(re.findall(r"^[E>] +(.*)$", pytest_error_string, re.MULTILINE))
def extract_unique_errors(pytest_output: str) -> set[str]:
unique_errors = set()
# Regex pattern to match error lines:
# - Start with 'E' followed by optional whitespace
# - Capture the actual error message
pattern = r"^E\s+(.*)$"
for error_message in re.findall(pattern, pytest_output, re.MULTILINE):
error_message = error_message.strip()
if error_message:
unique_errors.add(error_message)
return unique_errors
def validate_relative_directory_path(path: str) -> tuple[bool, str]:
"""Validate that a path is a safe relative directory path.
Prevents path traversal attacks and invalid paths.
Works cross-platform (Windows, Linux, macOS).
Args:
path: The path string to validate
Returns:
tuple[bool, str]: (is_valid, error_message)
- is_valid: True if path is valid, False otherwise
- error_message: Empty string if valid, error description if invalid
"""
if not path or not path.strip():
return False, "Path cannot be empty"
# Normalize whitespace
path = path.strip()
# Check for path traversal attempts (cross-platform)
# Normalize path separators for checking
normalized = path.replace("\\", "/")
if ".." in normalized:
return False, "Path cannot contain '..'. Use a relative path like 'tests' or 'src/app' instead"
# Check for absolute paths, invalid characters, and validate path format
error_msg = ""
if Path(path).is_absolute():
error_msg = "Path must be relative, not absolute"
elif os.name == "nt": # Windows
if any(char in _INVALID_CHARS_NT for char in path):
error_msg = "Path contains invalid characters for this operating system"
elif "\0" in path: # Unix-like
error_msg = "Path contains invalid characters for this operating system"
else:
# Validate using pathlib to ensure it's a valid path structure
try:
Path(path)
except (ValueError, OSError) as e:
error_msg = f"Invalid path format: {e!s}"
if error_msg:
return False, error_msg
return True, ""
You can’t perform that action at this time.