Skip to content

Commit fe64706

Browse files
⚡️ Speed up function get_diff_lines_count by 15% in PR #274 (skip-formatting-for-large-diffs)
Here is an optimized version of your program. Key Improvements. - Avoids splitting all lines and list allocation; instead, iterates only as needed and sums matches (saves both memory and runtime). - Eliminates the inner function and replaces it with a fast inline check. **Why this is faster:** - Uses a simple for-loop instead of building a list. - Checks first character directly—less overhead than calling `startswith` multiple times. - Skips the closure. - No intermediate list storage. The function result and behavior are identical.
1 parent a5343fd commit fe64706

File tree

1 file changed

+22
-15
lines changed

1 file changed

+22
-15
lines changed

codeflash/code_utils/formatter.py

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,22 @@
44
import shlex
55
import subprocess
66
from typing import TYPE_CHECKING, Optional
7+
78
import isort
89

910
from codeflash.cli_cmds.console import console, logger
1011

1112
if TYPE_CHECKING:
1213
from pathlib import Path
1314

15+
1416
def get_nth_line(text: str, n: int) -> str | None:
1517
for i, line in enumerate(text.splitlines(), start=1):
1618
if i == n:
1719
return line
1820
return None
1921

22+
2023
def get_diff_output(cmd: list[str]) -> Optional[str]:
2124
try:
2225
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
@@ -27,54 +30,58 @@ def get_diff_output(cmd: list[str]) -> Optional[str]:
2730
is_ruff = cmd[0] == "ruff"
2831
if e.returncode == 0 and is_ruff:
2932
return ""
30-
elif e.returncode == 1 and is_ruff:
33+
if e.returncode == 1 and is_ruff:
3134
return e.stdout.strip() or None
3235
return None
3336

3437

3538
def get_diff_lines_output_by_black(filepath: str) -> Optional[str]:
3639
try:
3740
import black # type: ignore
38-
return get_diff_output(['black', '--diff', filepath])
41+
42+
return get_diff_output(["black", "--diff", filepath])
3943
except ImportError:
4044
return None
4145

46+
4247
def get_diff_lines_output_by_ruff(filepath: str) -> Optional[str]:
4348
try:
4449
import ruff # type: ignore
45-
return get_diff_output(['ruff', 'format', '--diff', filepath])
50+
51+
return get_diff_output(["ruff", "format", "--diff", filepath])
4652
except ImportError:
4753
print("can't import ruff")
4854
return None
4955

5056

5157
def get_diff_lines_count(diff_output: str) -> int:
52-
lines = diff_output.split('\n')
53-
def is_diff_line(line: str) -> bool:
54-
return line.startswith(('+', '-')) and not line.startswith(('+++', '---'))
55-
diff_lines = [line for line in lines if is_diff_line(line)]
56-
return len(diff_lines)
58+
# Count lines that are diff changes (start with '+' or '-', but not '+++' or '---')
59+
count = 0
60+
for line in diff_output.split("\n"):
61+
if line and (line[0] in ("+", "-")) and not (line[:3] == "+++" or line[:3] == "---"):
62+
count += 1
63+
return count
64+
5765

5866
def is_safe_to_format(filepath: str, max_diff_lines: int = 100) -> bool:
5967
diff_changes_stdout = None
6068

6169
diff_changes_stdout = get_diff_lines_output_by_black(filepath)
6270

6371
if diff_changes_stdout is None:
64-
logger.warning(f"black formatter not found, trying ruff instead...")
72+
logger.warning("black formatter not found, trying ruff instead...")
6573
diff_changes_stdout = get_diff_lines_output_by_ruff(filepath)
6674
if diff_changes_stdout is None:
67-
logger.warning(f"Both ruff, black formatters not found, skipping formatting diff check.")
75+
logger.warning("Both ruff, black formatters not found, skipping formatting diff check.")
6876
return False
69-
77+
7078
diff_lines_count = get_diff_lines_count(diff_changes_stdout)
71-
79+
7280
if diff_lines_count > max_diff_lines:
7381
logger.debug(f"Skipping {filepath}: {diff_lines_count} lines would change (max: {max_diff_lines})")
7482
return False
75-
else:
76-
return True
77-
83+
return True
84+
7885

7986
def format_code(formatter_cmds: list[str], path: Path, print_status: bool = True) -> str: # noqa
8087
# TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution

0 commit comments

Comments
 (0)