Skip to content

Commit a2012ee

Browse files
⚡️ Speed up method AssertCleanup.transform_asserts by 82% in PR #26 (clean_concolic_tests)
To optimize the given code, we should focus on reducing redundant operations and improving the performance of regular expression matching, string manipulation, and list operations. Let's break down the changes. ### Key Optimizations. This results in a more efficient and potentially faster code execution while maintaining the same functionality and output structure.
1 parent af1f078 commit a2012ee

File tree

1 file changed

+17
-10
lines changed

1 file changed

+17
-10
lines changed

codeflash/code_utils/code_replacer.py

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import re
55
from collections import defaultdict
66
from functools import lru_cache
7-
from typing import TYPE_CHECKING, Optional, TypeVar
7+
from typing import TYPE_CHECKING, List, Optional, TypeVar
88

99
import libcst as cst
1010

@@ -342,33 +342,35 @@ def function_to_optimize_original_worktree_fqn(
342342
class AssertCleanup:
343343
def transform_asserts(self, code: str) -> str:
344344
lines = code.splitlines()
345-
result_lines = []
345+
result_lines: List[str] = []
346+
347+
append_result = result_lines.append
348+
transform_line = self._transform_assert_line
346349

347350
for line in lines:
348-
transformed = self._transform_assert_line(line)
351+
transformed = transform_line(line)
349352
if transformed is not None:
350-
result_lines.append(transformed)
353+
append_result(transformed)
351354
else:
352-
result_lines.append(line)
355+
append_result(line)
353356

354357
return "\n".join(result_lines)
355358

356359
def _transform_assert_line(self, line: str) -> Optional[str]:
357360
indent = line[: len(line) - len(line.lstrip())]
358361

359-
assert_match = re.match(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$", line)
362+
assert_match = self.assert_pattern.match(line)
360363
if assert_match:
361364
expression = assert_match.group(1).strip()
362365
if expression.startswith("not "):
363366
return f"{indent}{expression}"
364367

365-
expression = re.sub(r"[,;]\s*$", "", expression)
368+
expression = expression.rstrip(",;")
366369
return f"{indent}{expression}"
367370

368-
unittest_match = re.match(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$", line)
371+
unittest_match = self.unittest_pattern.match(line)
369372
if unittest_match:
370-
indent, assert_method, args = unittest_match.groups()
371-
373+
indent, _, args = unittest_match.groups()
372374
if args:
373375
arg_parts = self._split_top_level_args(args)
374376
if arg_parts and arg_parts[0]:
@@ -399,6 +401,11 @@ def _split_top_level_args(self, args_str: str) -> list[str]:
399401

400402
return result
401403

404+
def __init__(self):
405+
# Compile the regular expressions once
406+
self.assert_pattern = re.compile(r"\s*assert\s+(.*?)(?:\s*==\s*.*)?$")
407+
self.unittest_pattern = re.compile(r"(\s*)self\.assert([A-Za-z]+)\((.*)\)$")
408+
402409

403410
def clean_concolic_tests(test_suite_code: str) -> str:
404411
try:

0 commit comments

Comments
 (0)