refactor: reduce cognitive complexity across diffctx modules

Author: nikolay-e

src/treemapper/diffctx/__init__.py

@@ -316,6 +316,33 @@ def _validate_inputs(root_dir: Path, alpha: float, tau: float, budget_tokens: in
         raise ValueError(f"budget_tokens must be > 0, got {budget_tokens}")
 
 
+def _find_dangling_semantic_names(
+    selected: list[Fragment],
+    graph: Graph,
+    frag_by_id: dict[FragmentId, Fragment],
+    selected_ids: set[FragmentId],
+) -> set[str]:
+    dangling: set[str] = set()
+    for frag in selected:
+        for nbr_id in graph.neighbors(frag.id):
+            if nbr_id in selected_ids:
+                continue
+            if graph.edge_categories.get((frag.id, nbr_id), "") != "semantic":
+                continue
+            nbr_frag = frag_by_id.get(nbr_id)
+            if nbr_frag and nbr_frag.symbol_name:
+                dangling.add(nbr_frag.symbol_name.lower())
+    return dangling
+
+
+def _pick_best_fragment(candidates: list[Fragment], selected_ids: set[FragmentId]) -> Fragment | None:
+    if any(c.id in selected_ids for c in candidates):
+        return None
+    sig_candidates = [f for f in candidates if "_signature" in f.kind]
+    full_candidates = [f for f in candidates if "_signature" not in f.kind]
+    return next(iter(sig_candidates or full_candidates), None)
+
+
 def _coherence_post_pass(
     result: SelectionResult,
     all_fragments: list[Fragment],
@@ -331,32 +358,15 @@ def _coherence_post_pass(
             name_to_frags.setdefault(f.symbol_name.lower(), []).append(f)
 
     frag_by_id: dict[FragmentId, Fragment] = {f.id: f for f in all_fragments}
-
-    dangling_names: set[str] = set()
-    for frag in result.selected:
-        for nbr_id in graph.neighbors(frag.id):
-            if nbr_id in selected_ids:
-                continue
-            cat = graph.edge_categories.get((frag.id, nbr_id), "")
-            if cat == "semantic":
-                nbr_frag = frag_by_id.get(nbr_id)
-                if nbr_frag and nbr_frag.symbol_name:
-                    dangling_names.add(nbr_frag.symbol_name.lower())
+    dangling_names = _find_dangling_semantic_names(result.selected, graph, frag_by_id, selected_ids)
 
     added: list[Fragment] = []
     for name in dangling_names:
-        candidates = name_to_frags.get(name, [])
-        for c in candidates:
-            if c.id in selected_ids:
-                break
-        else:
-            sig_candidates = [f for f in candidates if "_signature" in f.kind]
-            full_candidates = [f for f in candidates if "_signature" not in f.kind]
-            pick = next(iter(sig_candidates or full_candidates), None)
-            if pick and pick.token_count <= remaining and pick.id not in selected_ids:
-                added.append(pick)
-                selected_ids.add(pick.id)
-                remaining -= pick.token_count
+        pick = _pick_best_fragment(name_to_frags.get(name, []), selected_ids)
+        if pick and pick.token_count <= remaining and pick.id not in selected_ids:
+            added.append(pick)
+            selected_ids.add(pick.id)
+            remaining -= pick.token_count
 
     if not added:
         return result

src/treemapper/diffctx/edges/base.py

@@ -22,26 +22,31 @@
 _INDEX_FILE_STEMS = frozenset({"__init__", "index", "mod"})
 
 
+def _strip_source_prefix(parts: list[str]) -> list[str]:
+    for i, part in enumerate(parts):
+        if part in ("src", "lib", "packages"):
+            return parts[i + 1 :]
+    return parts
+
+
+def _strip_file_extension(stem: str) -> str:
+    for ext in sorted(_STRIP_EXTENSIONS, key=len, reverse=True):
+        if stem.endswith(ext):
+            return stem[: -len(ext)]
+    return stem
+
+
 def path_to_module(path: Path, repo_root: Path | None = None) -> str:
     if repo_root and path.is_absolute():
         try:
             path = path.relative_to(repo_root)
         except ValueError:
             pass
 
-    parts = list(path.parts)
-
-    for i, part in enumerate(parts):
-        if part in ("src", "lib", "packages"):
-            parts = parts[i + 1 :]
-            break
+    parts = _strip_source_prefix(list(path.parts))
 
     if parts:
-        stem = parts[-1]
-        for ext in sorted(_STRIP_EXTENSIONS, key=len, reverse=True):
-            if stem.endswith(ext):
-                parts[-1] = stem[: -len(ext)]
-                break
+        parts[-1] = _strip_file_extension(parts[-1])
         if parts and parts[-1] in _INDEX_FILE_STEMS:
             parts = parts[:-1]
 

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

@@ -26,7 +26,7 @@
     re.MULTILINE,
 )
 _KOTLIN_FUN_RE = re.compile(
-    r"^\s*(?:(?:public|private|internal|protected|abstract|open|final|override|inline|external|tailrec|operator|infix|suspend|expect|actual)\s+)*fun\s+(?:<[^>]+>\s+)?([a-z]\w*)",
+    r"^\s*(?:\w+\s+)*fun\s+(?:<[^>]+>\s+)?([a-z]\w*)",
     re.MULTILINE,
 )
 
@@ -38,9 +38,7 @@
 _SCALA_DEF_RE = re.compile(r"^\s*(?:private |protected )?def\s+([a-z]\w*)", re.MULTILINE)
 
 _TYPE_REF_RE = re.compile(r"(?<![a-z_])([A-Z]\w*)\b")
-_KOTLIN_INHERIT_RE = re.compile(
-    r"(?:class|interface|object)\s+\w+(?:<[^>]*>)?(?:\([^)]*\))?\s*:\s*([A-Z]\w*(?:\s*,\s*[A-Z]\w*)*)"
-)
+_KOTLIN_INHERIT_RE = re.compile(r"(?:class|interface|object)\s+\w+[^:\n]*:\s*([A-Z]\w*(?:\s*,\s*[A-Z]\w*)*)")
 _SCALA_WITH_RE = re.compile(r"\bwith\s+([A-Z]\w*)")
 
 _ANNOTATION_RE = re.compile(r"@([A-Z]\w*)")
@@ -88,27 +86,25 @@ def _extract_classes(content: str, path: Path) -> set[str]:
     return classes
 
 
-def _extract_inheritance(content: str, path: Path) -> set[str]:
+def _split_class_list(regex: re.Pattern[str], content: str) -> set[str]:
     refs: set[str] = set()
+    for m in regex.finditer(content):
+        for cls in m.group(1).split(","):
+            stripped = cls.strip()
+            if stripped:
+                refs.add(stripped)
+    return refs
+
+
+def _extract_inheritance(content: str, path: Path) -> set[str]:
     if _is_kotlin(path):
-        for m in _KOTLIN_INHERIT_RE.finditer(content):
-            for cls in m.group(1).split(","):
-                cls = cls.strip()
-                if cls:
-                    refs.add(cls)
-    elif _is_scala(path):
-        for m in _JAVA_EXTENDS_RE.finditer(content):
-            for cls in m.group(1).split(","):
-                refs.add(cls.strip())
-        for m in _SCALA_WITH_RE.finditer(content):
-            refs.add(m.group(1))
-    else:
-        for m in _JAVA_EXTENDS_RE.finditer(content):
-            for cls in m.group(1).split(","):
-                refs.add(cls.strip())
-        for m in _JAVA_IMPLEMENTS_RE.finditer(content):
-            for cls in m.group(1).split(","):
-                refs.add(cls.strip())
+        return _split_class_list(_KOTLIN_INHERIT_RE, content)
+    if _is_scala(path):
+        refs = _split_class_list(_JAVA_EXTENDS_RE, content)
+        refs.update(m.group(1) for m in _SCALA_WITH_RE.finditer(content))
+        return refs
+    refs = _split_class_list(_JAVA_EXTENDS_RE, content)
+    refs.update(_split_class_list(_JAVA_IMPLEMENTS_RE, content))
     return refs
 
 

src/treemapper/diffctx/edges/structural/test.py

@@ -100,32 +100,34 @@ def discover_related_files(
         repo_root: Path | None = None,
     ) -> list[Path]:
         changed_set = set(changed_files)
-        discovered: list[Path] = []
-
         candidate_by_stem: dict[str, list[Path]] = defaultdict(list)
         for c in all_candidate_files:
             if c not in changed_set:
                 candidate_by_stem[c.stem.lower()].append(c)
 
+        discovered: list[Path] = []
         for changed in changed_files:
-            stem = changed.stem.lower()
             suffix = changed.suffix.lower()
-
             if _is_test_file(changed):
-                target = _extract_target_name_from_test(changed.stem)
-                if target:
-                    for cand in candidate_by_stem.get(target, []):
-                        if cand.suffix.lower() == suffix:
-                            discovered.append(cand)
+                discovered.extend(self._find_source_for_test(changed, suffix, candidate_by_stem))
             else:
-                test_stems = [f"test_{stem}", f"{stem}_test"]
-                for ts in test_stems:
-                    for cand in candidate_by_stem.get(ts, []):
-                        if cand.suffix.lower() == suffix and _is_test_file(cand):
-                            discovered.append(cand)
-
+                discovered.extend(self._find_tests_for_source(changed.stem.lower(), suffix, candidate_by_stem))
         return discovered
 
+    @staticmethod
+    def _find_source_for_test(changed: Path, suffix: str, candidate_by_stem: dict[str, list[Path]]) -> list[Path]:
+        target = _extract_target_name_from_test(changed.stem)
+        if not target:
+            return []
+        return [c for c in candidate_by_stem.get(target, []) if c.suffix.lower() == suffix]
+
+    @staticmethod
+    def _find_tests_for_source(stem: str, suffix: str, candidate_by_stem: dict[str, list[Path]]) -> list[Path]:
+        results: list[Path] = []
+        for ts in (f"test_{stem}", f"{stem}_test"):
+            results.extend(c for c in candidate_by_stem.get(ts, []) if c.suffix.lower() == suffix and _is_test_file(c))
+        return results
+
     def build(self, fragments: list[Fragment], repo_root: Path | None = None) -> EdgeDict:
         edges: EdgeDict = {}
 

src/treemapper/diffctx/parsers/markdown.py

@@ -63,28 +63,42 @@ def _find_all_headings(self, lines: list[str]) -> list[tuple[int, int]]:
         headings: list[tuple[int, int]] = []
         i = 0
         while i < len(lines):
-            stripped = lines[i].lstrip()
-            if stripped.startswith("#"):
-                level = len(stripped) - len(stripped.lstrip("#"))
-                if level <= 6 and (len(stripped) == level or stripped[level] == " "):
-                    headings.append((i + 1, level))
+            result = self._try_atx_heading(lines[i])
+            if result is not None:
+                headings.append((i + 1, result))
                 i += 1
                 continue
 
-            if stripped and i + 1 < len(lines):
-                next_stripped = lines[i + 1].strip()
-                if next_stripped and all(c == "=" for c in next_stripped):
-                    headings.append((i + 1, 1))
-                    i += 2
-                    continue
-                if len(next_stripped) >= 2 and all(c == "-" for c in next_stripped):
-                    headings.append((i + 1, 2))
-                    i += 2
-                    continue
+            setext = self._try_setext_heading(lines, i)
+            if setext is not None:
+                headings.append((i + 1, setext))
+                i += 2
+                continue
 
             i += 1
         return headings
 
+    @staticmethod
+    def _try_atx_heading(line: str) -> int | None:
+        stripped = line.lstrip()
+        if not stripped.startswith("#"):
+            return None
+        level = len(stripped) - len(stripped.lstrip("#"))
+        if level <= 6 and (len(stripped) == level or stripped[level] == " "):
+            return level
+        return None
+
+    @staticmethod
+    def _try_setext_heading(lines: list[str], i: int) -> int | None:
+        if not lines[i].lstrip() or i + 1 >= len(lines):
+            return None
+        next_stripped = lines[i + 1].strip()
+        if next_stripped and all(c == "=" for c in next_stripped):
+            return 1
+        if len(next_stripped) >= 2 and all(c == "-" for c in next_stripped):
+            return 2
+        return None
+
     def _find_section_end(self, lines: list[str], _start_line: int, level: int, remaining_headings: list[tuple[int, int]]) -> int:
         for next_line, next_level in remaining_headings:
             if next_level <= level: