diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py index a9050c7ca..8a6811675 100644 --- a/codeflash/languages/java/remove_asserts.py +++ b/codeflash/languages/java/remove_asserts.py @@ -192,6 +192,7 @@ def __init__( # Precompile the assignment-detection regex to avoid recompiling on each call. self._assign_re = re.compile(r"(\w+(?:<[^>]+>)?)\s+(\w+)\s*=\s*$") + self._special_re = re.compile(r"""['"()]""") def transform(self, source: str) -> str: """Remove assertions from source code, preserving target function calls. @@ -804,17 +805,20 @@ def _find_balanced_parens(self, code: str, open_paren_pos: int) -> tuple[str | N string_char = None in_char = False - # Track previous character locally to avoid repeated indexing (code[pos-1]). - prev_char = code[open_paren_pos] + while depth > 0: + m = self._special_re.search(code, pos) + if m is None: + return None, -1 - while pos < end and depth > 0: - char = code[pos] + i = m.start() + char = m.group() + escaped = i > 0 and code[i - 1] == "\\" # Handle character literals - if char == "'" and not in_string and prev_char != "\\": + if char == "'" and not in_string and not escaped: in_char = not in_char # Handle string literals (double quotes) - elif char == '"' and not in_char and prev_char != "\\": + elif char == '"' and not in_char and not escaped: if not in_string: in_string = True string_char = char @@ -827,13 +831,7 @@ def _find_balanced_parens(self, code: str, open_paren_pos: int) -> tuple[str | N elif char == ")": depth -= 1 - pos += 1 - - prev_char = char - - if depth != 0: - return None, -1 - + pos = i + 1 return code[open_paren_pos + 1 : pos - 1], pos def _find_balanced_braces(self, code: str, open_brace_pos: int) -> tuple[str | None, int]: diff --git a/codeflash/languages/registry.py b/codeflash/languages/registry.py index 38688cab6..e32bb5c16 100644 --- a/codeflash/languages/registry.py +++ b/codeflash/languages/registry.py @@ -56,7 +56,7 @@ def _ensure_languages_registered() -> None: from codeflash.languages.javascript import support as _ with contextlib.suppress(ImportError): - from codeflash.languages.java import support as _ + from codeflash.languages.java import support as _ # noqa: F401 _languages_registered = True