fix: resolve all SonarCloud issues — ReDoS, complexity, regex style

Author: nikolay-e

src/treemapper/diffctx/__init__.py

@@ -214,6 +214,30 @@ def _truncate_generated_fragments(file_frags: list[Fragment]) -> list[Fragment]:
     return truncated
 
 
+def _dedup_fragments(raw_frags: list[Fragment]) -> list[Fragment]:
+    seen: set[FragmentId] = set()
+    result: list[Fragment] = []
+    for f in raw_frags:
+        if f.id not in seen:
+            result.append(f)
+            seen.add(f.id)
+    return result
+
+
+def _cap_fragments(file_frags: list[Fragment], cap: int, file_path: Path, is_generated: bool) -> list[Fragment]:
+    if len(file_frags) <= cap:
+        return file_frags
+    file_frags.sort(key=lambda f: f.line_count, reverse=True)
+    file_frags = file_frags[:cap]
+    logger.debug(
+        "diffctx: capped %s to %d fragments%s",
+        file_path.name,
+        cap,
+        " (generated)" if is_generated else "",
+    )
+    return file_frags
+
+
 def _process_files_for_fragments(
     files: list[Path],
     root_dir: Path,
@@ -227,25 +251,11 @@ def _process_files_for_fragments(
         if content is None:
             continue
         raw_frags = [f for f in fragment_file(file_path, content) if f.id not in seen_frag_ids]
-        local_seen: set[FragmentId] = set()
-        file_frags = []
-        for f in raw_frags:
-            if f.id not in local_seen:
-                file_frags.append(f)
-                local_seen.add(f.id)
+        file_frags = _dedup_fragments(raw_frags)
 
         is_generated = _is_generated_file(file_path, content)
         cap = _MAX_GENERATED_FRAGMENTS if is_generated else max_frags
-
-        if len(file_frags) > cap:
-            file_frags.sort(key=lambda f: f.line_count, reverse=True)
-            file_frags = file_frags[:cap]
-            logger.debug(
-                "diffctx: capped %s to %d fragments%s",
-                file_path.name,
-                cap,
-                " (generated)" if is_generated else "",
-            )
+        file_frags = _cap_fragments(file_frags, cap, file_path, is_generated)
 
         if is_generated:
             file_frags = _truncate_generated_fragments(file_frags)

src/treemapper/diffctx/edges/semantic/go.py

@@ -113,6 +113,31 @@ def _get_package_name_from_content(content: str, path: Path) -> str:
     return path.parent.name
 
 
+def _resolve_bases(pattern_str: str, parent: Path, repo_root: Path | None) -> list[Path]:
+    base_pattern = pattern_str.split("*")[0].rstrip("/")
+    candidate_bases = [parent / base_pattern]
+    if repo_root:
+        candidate_bases.append(repo_root / base_pattern)
+    dirs: list[Path] = []
+    for base in candidate_bases:
+        try:
+            dirs.append(base.resolve())
+        except (OSError, ValueError):
+            pass
+    return dirs
+
+
+def _any_dir_matches(dirs_to_check: set[Path], embed_dirs: list[Path]) -> bool:
+    for d in dirs_to_check:
+        try:
+            resolved = d.resolve()
+            if any(resolved == ed or resolved.is_relative_to(ed) for ed in embed_dirs):
+                return True
+        except (ValueError, OSError):
+            continue
+    return False
+
+
 class GoEdgeBuilder(EdgeBuilder):
     weight = 0.75
     import_weight = EDGE_WEIGHTS["go_import"].forward
@@ -127,55 +152,55 @@ def discover_related_files(
         all_candidate_files: list[Path],
         repo_root: Path | None = None,
     ) -> list[Path]:
-        go_changed = [f for f in changed_files if _is_go_file(f)]
         changed_set = set(changed_files)
         discovered: set[Path] = set()
+        candidates = [c for c in all_candidate_files if c not in changed_set and _is_go_file(c)]
 
+        go_changed = [f for f in changed_files if _is_go_file(f)]
         if go_changed:
-            changed_pkg_dirs = {f.parent for f in go_changed}
-            for candidate in all_candidate_files:
-                if candidate not in changed_set and _is_go_file(candidate) and candidate.parent in changed_pkg_dirs:
-                    discovered.add(candidate)
+            self._discover_same_package(go_changed, candidates, discovered)
+
+        embed_go_files = self._discover_embed_files(changed_files, candidates, discovered, repo_root)
+        self._discover_package_peers(embed_go_files, candidates, discovered)
+
+        return list(discovered)
+
+    def _discover_same_package(self, go_changed: list[Path], candidates: list[Path], discovered: set[Path]) -> None:
+        pkg_dirs = {f.parent for f in go_changed}
+        for c in candidates:
+            if c.parent in pkg_dirs:
+                discovered.add(c)
 
+    def _discover_embed_files(
+        self,
+        changed_files: list[Path],
+        candidates: list[Path],
+        discovered: set[Path],
+        repo_root: Path | None,
+    ) -> set[Path]:
         changed_dirs = {f.parent for f in changed_files}
         embed_go_files: set[Path] = set()
-        for candidate in all_candidate_files:
-            if candidate not in changed_set and _is_go_file(candidate):
-                if self._embeds_any_changed_dir(candidate, changed_dirs, repo_root):
-                    discovered.add(candidate)
-                    embed_go_files.add(candidate)
+        for c in candidates:
+            if self._embeds_any_changed_dir(c, changed_dirs, repo_root):
+                discovered.add(c)
+                embed_go_files.add(c)
+        return embed_go_files
 
+    def _discover_package_peers(self, embed_go_files: set[Path], candidates: list[Path], discovered: set[Path]) -> None:
         embed_dirs = {f.parent for f in embed_go_files}
-        for candidate in all_candidate_files:
-            if candidate not in changed_set and _is_go_file(candidate) and candidate not in discovered:
-                if candidate.parent in embed_dirs:
-                    discovered.add(candidate)
-
-        return list(discovered)
+        for c in candidates:
+            if c not in discovered and c.parent in embed_dirs:
+                discovered.add(c)
 
     def _embeds_any_changed_dir(self, go_file: Path, changed_dirs: set[Path], repo_root: Path | None = None) -> bool:
         try:
             content = go_file.read_text(encoding="utf-8")
         except (OSError, UnicodeDecodeError):
             return False
         for match in _GO_EMBED_RE.finditer(content):
-            embed_pattern = match.group(1)
-            base_pattern = embed_pattern.split("*")[0].rstrip("/")
-            candidate_bases = [go_file.parent / base_pattern]
-            if repo_root:
-                candidate_bases.append(repo_root / base_pattern)
-            for base in candidate_bases:
-                try:
-                    embed_dir = base.resolve()
-                except (OSError, ValueError):
-                    continue
-                for changed_dir in changed_dirs:
-                    try:
-                        resolved = changed_dir.resolve()
-                        if resolved == embed_dir or resolved.is_relative_to(embed_dir):
-                            return True
-                    except (ValueError, OSError):
-                        continue
+            embed_dirs = _resolve_bases(match.group(1), go_file.parent, repo_root)
+            if _any_dir_matches(changed_dirs, embed_dirs):
+                return True
         return False
 
     def build(self, fragments: list[Fragment], repo_root: Path | None = None) -> EdgeDict:
@@ -199,28 +224,26 @@ def _build_embed_edges(
         edges: EdgeDict,
         repo_root: Path | None = None,
     ) -> None:
+        non_go_frags = [f for f in all_frags if not _is_go_file(f.path)]
         for gf in go_frags:
             for match in _GO_EMBED_RE.finditer(gf.content):
-                embed_pattern = match.group(1)
-                base_pattern = embed_pattern.split("*")[0].rstrip("/")
-                candidate_bases = [gf.path.parent / base_pattern]
-                if repo_root:
-                    candidate_bases.append(repo_root / base_pattern)
-                embed_dirs: list[Path] = []
-                for base in candidate_bases:
-                    try:
-                        embed_dirs.append(base.resolve())
-                    except (OSError, ValueError):
-                        pass
-                for frag in all_frags:
-                    if _is_go_file(frag.path):
-                        continue
-                    try:
-                        frag_resolved = frag.path.resolve()
-                        if any(frag_resolved.is_relative_to(ed) for ed in embed_dirs):
-                            self.add_edge(edges, gf.id, frag.id, self.weight * 0.8)
-                    except (ValueError, OSError):
-                        continue
+                embed_dirs = _resolve_bases(match.group(1), gf.path.parent, repo_root)
+                self._link_embed_targets(gf, non_go_frags, embed_dirs, edges)
+
+    def _link_embed_targets(
+        self,
+        gf: Fragment,
+        non_go_frags: list[Fragment],
+        embed_dirs: list[Path],
+        edges: EdgeDict,
+    ) -> None:
+        for frag in non_go_frags:
+            try:
+                frag_resolved = frag.path.resolve()
+                if any(frag_resolved.is_relative_to(ed) for ed in embed_dirs):
+                    self.add_edge(edges, gf.id, frag.id, self.weight * 0.8)
+            except (ValueError, OSError):
+                continue
 
     def _build_indices(
         self, go_frags: list[Fragment], repo_root: Path | None

src/treemapper/diffctx/edges/semantic/javascript.py

@@ -286,28 +286,39 @@ def _add_import_edges(
             file_to_frags[f.path].append(f.id)
 
         fragment_paths = set(file_to_frags.keys())
-
-        file_imports: dict[Path, set[str]] = defaultdict(set)
-        for f in js_frags:
-            for import_source in info_cache[f.id].imports:
-                if import_source.startswith("."):
-                    file_imports[f.path].add(import_source)
+        file_imports = self._collect_relative_imports(js_frags, info_cache)
 
         for src_path, import_sources in file_imports.items():
             for import_source in import_sources:
                 resolved = _resolve_relative_import(src_path, import_source, fragment_paths)
                 if resolved is None or resolved == src_path:
                     continue
                 target_ids = file_to_frags.get(resolved, [])
-                if not target_ids:
-                    continue
-                for src_id in file_to_frags[src_path]:
-                    for target_id in target_ids:
-                        if target_id == src_id:
-                            continue
-                        w = self._IMPORT_WEIGHT
-                        edges[(src_id, target_id)] = max(edges.get((src_id, target_id), 0.0), w)
-                        edges[(target_id, src_id)] = max(edges.get((target_id, src_id), 0.0), w * self.reverse_weight_factor)
+                if target_ids:
+                    self._link_import_pairs(file_to_frags[src_path], target_ids, edges)
+
+    @staticmethod
+    def _collect_relative_imports(js_frags: list[Fragment], info_cache: dict[FragmentId, JsFragmentInfo]) -> dict[Path, set[str]]:
+        file_imports: dict[Path, set[str]] = defaultdict(set)
+        for f in js_frags:
+            for import_source in info_cache[f.id].imports:
+                if import_source.startswith("."):
+                    file_imports[f.path].add(import_source)
+        return file_imports
+
+    def _link_import_pairs(
+        self,
+        src_ids: list[FragmentId],
+        target_ids: list[FragmentId],
+        edges: EdgeDict,
+    ) -> None:
+        w = self._IMPORT_WEIGHT
+        rev_w = w * self.reverse_weight_factor
+        for src_id in src_ids:
+            for target_id in target_ids:
+                if target_id != src_id:
+                    edges[(src_id, target_id)] = max(edges.get((src_id, target_id), 0.0), w)
+                    edges[(target_id, src_id)] = max(edges.get((target_id, src_id), 0.0), rev_w)
 
     def _discover_forward_imports(
         self,

src/treemapper/diffctx/edges/semantic/php.py

@@ -10,8 +10,8 @@
 
 _PHP_EXTS = {".php", ".phtml", ".php3", ".php4", ".php5", ".php7", ".phps"}
 
-_PHP_USE_RE = re.compile(r"^\s*use\s+([A-Z][a-zA-Z0-9_]*\\[a-zA-Z0-9_\\]*(?:\s+as\s+\w+)?)\s*;", re.MULTILINE)
-_PHP_NAMESPACE_RE = re.compile(r"^\s*namespace\s+([A-Z][a-zA-Z0-9_\\]*)\s*;", re.MULTILINE)
+_PHP_USE_RE = re.compile(r"^\s*use\s+([A-Z]\w*\\[\w\\]*(?:\s+as\s+\w+)?)\s*;", re.MULTILINE)
+_PHP_NAMESPACE_RE = re.compile(r"^\s*namespace\s+([A-Z][\w\\]*)\s*;", re.MULTILINE)
 _PHP_REQUIRE_RE = re.compile(r"^\s*(?:require|require_once|include|include_once)\s*\(?['\"]([^'\"]+)['\"]", re.MULTILINE)
 
 _PHP_CLASS_RE = re.compile(r"^\s*(?:abstract\s+)?(?:final\s+)?class\s+([A-Z]\w*)", re.MULTILINE)

src/treemapper/diffctx/edges/semantic/rust.py

@@ -16,7 +16,7 @@
 _RUST_STRUCT_RE = re.compile(r"^\s*(?:pub(?:\([^)]*\))?\s+)?struct\s+([A-Z]\w*)", re.MULTILINE)
 _RUST_ENUM_RE = re.compile(r"^\s*(?:pub(?:\([^)]*\))?\s+)?enum\s+([A-Z]\w*)", re.MULTILINE)
 _RUST_TRAIT_RE = re.compile(r"^\s*(?:pub(?:\([^)]*\))?\s+)?trait\s+([A-Z]\w*)", re.MULTILINE)
-_RUST_IMPL_RE = re.compile(r"^\s*impl(?:<[^<>]*(?:<[^<>]*>)*[^<>]*>)?\s+(?:\w+\s+for\s+)?([A-Z]\w*)", re.MULTILINE)
+_RUST_IMPL_RE = re.compile(r"^\s*impl(?:<[^>\n]*>)?\s+(?:\w+\s+for\s+)?([A-Z]\w*)", re.MULTILINE)
 _RUST_TYPE_ALIAS_RE = re.compile(r"^\s*(?:pub(?:\([^)]*\))?\s+)?type\s+([A-Z]\w*)", re.MULTILINE)
 
 _RUST_TYPE_REF_RE = re.compile(r"(?<![a-z_])([A-Z]\w*)\b")

src/treemapper/diffctx/parsers/base.py

@@ -104,35 +104,23 @@ def _find_first_balanced_point(lines: list[str], start_idx: int, target_end_idx:
     return None
 
 
-def find_balanced_end_line(
-    lines: list[str], start_idx: int, target_end_idx: int, max_extension: int = _GENERIC_MAX_EXTENSION
-) -> int:
-    if target_end_idx >= len(lines):
-        target_end_idx = len(lines) - 1
-
-    state = _BracketState()
-    target_state: _BracketState | None = None
-    for idx in range(start_idx, target_end_idx + 1):
+def _feed_lines(state: _BracketState, lines: list[str], start_idx: int, end_idx: int) -> None:
+    for idx in range(start_idx, end_idx + 1):
         if idx > start_idx:
             state.feed("\n")
         state.feed(lines[idx])
-    target_state = state
-
-    if target_state.depth == 0:
-        first_balanced = _find_first_balanced_point(lines, start_idx, target_end_idx)
-        if first_balanced is not None and first_balanced < target_end_idx:
-            return first_balanced
-        return target_end_idx
 
-    max_end = min(len(lines) - 1, target_end_idx + max_extension)
 
-    state = target_state.copy()
-    for end_idx in range(target_end_idx + 1, max_end + 1):
+def _scan_forward_for_balance(state: _BracketState, lines: list[str], start: int, end: int) -> int | None:
+    for end_idx in range(start, end + 1):
         state.feed("\n")
         state.feed(lines[end_idx])
         if state.depth == 0:
             return end_idx
+    return None
+
 
+def _scan_backward_for_balance(lines: list[str], start_idx: int, target_end_idx: int) -> int | None:
     state = _BracketState()
     last_balanced_idx = None
     for idx in range(start_idx, target_end_idx):
@@ -141,11 +129,30 @@ def find_balanced_end_line(
         state.feed(lines[idx])
         if state.depth == 0:
             last_balanced_idx = idx
+    return last_balanced_idx
 
-    if last_balanced_idx is not None:
-        return last_balanced_idx
 
-    return target_end_idx
+def find_balanced_end_line(
+    lines: list[str], start_idx: int, target_end_idx: int, max_extension: int = _GENERIC_MAX_EXTENSION
+) -> int:
+    if target_end_idx >= len(lines):
+        target_end_idx = len(lines) - 1
+
+    target_state = _BracketState()
+    _feed_lines(target_state, lines, start_idx, target_end_idx)
+
+    if target_state.depth == 0:
+        first_balanced = _find_first_balanced_point(lines, start_idx, target_end_idx)
+        if first_balanced is not None and first_balanced < target_end_idx:
+            return first_balanced
+        return target_end_idx
+
+    max_end = min(len(lines) - 1, target_end_idx + max_extension)
+    forward = _scan_forward_for_balance(target_state.copy(), lines, target_end_idx + 1, max_end)
+    if forward is not None:
+        return forward
+
+    return _scan_backward_for_balance(lines, start_idx, target_end_idx) or target_end_idx
 
 
 _SENTENCE_ENDINGS = (".", "?", "!", '."', '?"', '!"', ".'", "?'", "!'")