loom-code 0.1.1__py3-none-any.whl

loom_code/loominit/extractor.py

@@ -0,0 +1,332 @@
+"""Top-level orchestrator — ``build_index(repo_root) -> LoomIndex``.
+
+This is the only public entry point for slice 1. It composes every
+private helper in the package into one deterministic pipeline:
+
+1. :func:`_files.discover_files` — walk the repo, hash every file,
+   capture mtime + git heat + ``is_test`` flag.
+2. :func:`_ast_walk.walk_python_file` — per-file AST extraction
+   producing ``(_RawSymbol, _RawImport, _RawDecorator)`` triples.
+3. :func:`_resolve.build_module_index` + :func:`resolve_imports` —
+   turn dotted module names into rel-path edges.
+4. :func:`_resolve.detect_api_surface` — flag files reachable from
+   ``__init__.py`` re-exports / ``__all__``.
+5. :func:`_graph.pagerank_file_graph` — file-level centrality from
+   the resolved import graph.
+6. :func:`_tests_map.build_test_map` — grep test files for each
+   symbol's bare name, emit citations.
+7. :func:`_resolve.extract_entry_points` — pyproject scripts,
+   ``__main__`` blocks, landmark-decorated callables.
+8. :func:`_graph.cluster_by_path_prefix` — group files into
+   subsystem clusters with hash buckets for diff-aware refresh.
+
+The result is a :class:`schema.LoomIndex`, consumed in-memory by
+:mod:`repomap` to render the per-turn repo map. No LLM calls anywhere
+here — the structural pass is deterministic by design.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Iterable
+from datetime import UTC, datetime
+from pathlib import Path
+
+from ._ast_walk import _RawDecorator, _RawSymbol, walk_python_file
+from ._files import DiscoveredFile, discover_files, is_git_repo
+from ._graph import cluster_by_path_prefix, pagerank_file_graph
+from ._resolve import (
+    build_module_index,
+    detect_api_surface,
+    extract_entry_points,
+    resolve_imports,
+)
+from ._tests_map import build_test_map
+from .schema import (
+    CallEdge,
+    Cluster,
+    DecoratorLandmark,
+    FileEntry,
+    LoomIndex,
+    SymbolEntry,
+)
+
+
+def build_index(repo_root: Path) -> LoomIndex:
+    """Run the full structural extraction pass and return the index.
+
+    Side effect: NONE — returns an in-memory :class:`LoomIndex` that
+    :mod:`repomap` renders into the per-turn repo map. Keeping it
+    I/O-free makes the pipeline trivially testable against in-memory
+    fixtures.
+    """
+    files = discover_files(repo_root)
+    py_files = [f for f in files if f.lang == "python"]
+
+    # ---- 1. Walk every Python file ------------------------------------
+    per_file = _walk_all(py_files)
+
+    # ---- 2. Module-name resolution / API surface ----------------------
+    module_index = build_module_index(files)
+    api_surface = detect_api_surface(files, module_index)
+
+    # ---- 3. Aggregate symbols + assign file-level PageRank ------------
+    raw_imports_by_file: dict[str, list[tuple[str, int, int]]] = {
+        path: [(imp.to_module, imp.line, imp.level) for imp in imps]
+        for path, (_syms, imps, _decs) in per_file.items()
+    }
+    import_edges = resolve_imports(raw_imports_by_file, module_index)
+
+    # File-level PageRank: edges restricted to resolved (in-repo) ones.
+    edges_for_pagerank = _pagerank_edges(
+        import_edges, py_files, module_index, raw_imports_by_file
+    )
+    file_scores = pagerank_file_graph(
+        files=[f.rel_path for f in py_files],
+        edges=edges_for_pagerank,
+    )
+
+    # n_callers / n_callees per file (resolved edges only).
+    in_degree, out_degree = _degree_maps(edges_for_pagerank)
+
+    # ---- 4. Test→symbol map -------------------------------------------
+    all_symbol_names = {
+        sym.name
+        for path, (syms, _imps, _decs) in per_file.items()
+        for sym in syms
+        if sym.is_public
+    }
+    test_map = build_test_map(files=files, symbol_names=all_symbol_names)
+
+    # ---- 5. Build schema.FileEntry list -------------------------------
+    file_entries = [
+        FileEntry(
+            path=f.rel_path,
+            lang=f.lang,
+            size_bytes=f.size_bytes,
+            lines=f.lines,
+            sha256=f.sha256,
+            mtime=f.mtime,
+            git_changes_90d=f.git_changes_90d,
+            is_test=f.is_test,
+            in_api_surface=f.rel_path in api_surface,
+        )
+        for f in files
+    ]
+
+    # ---- 6. Build schema.SymbolEntry list -----------------------------
+    symbol_entries: list[SymbolEntry] = []
+    for path, (raw_syms, _imps, _decs) in per_file.items():
+        file_pr = file_scores.get(path, 0.0)
+        file_in = in_degree.get(path, 0)
+        file_out = out_degree.get(path, 0)
+        for raw in raw_syms:
+            sym_id = f"{path}:{raw.qualified_name}"
+            symbol_entries.append(
+                SymbolEntry(
+                    id=sym_id,
+                    name=raw.name,
+                    qualified_name=raw.qualified_name,
+                    kind=raw.kind,
+                    path=path,
+                    line=raw.line,
+                    end_line=raw.end_line,
+                    signature=raw.signature,
+                    docstring_first_line=raw.docstring_first_line,
+                    decorators=list(raw.decorators),
+                    is_public=raw.is_public,
+                    in_api_surface=(path in api_surface) and raw.is_public,
+                    # File-level for now; symbol-level requires call graph
+                    pagerank=file_pr,
+                    n_callers=file_in,
+                    n_callees=file_out,
+                    tests=test_map.get(raw.name, []) if raw.is_public else [],
+                )
+            )
+
+    # ---- 7. Decorator landmarks --------------------------------------
+    decorator_entries: list[DecoratorLandmark] = []
+    decorator_path_lookup: dict[_RawDecorator, str] = {}
+    for path, (_syms, _imps, raw_decs) in per_file.items():
+        for raw in raw_decs:
+            decorator_path_lookup[raw] = path
+            decorator_entries.append(
+                DecoratorLandmark(
+                    decorator=raw.decorator,
+                    target=f"{path}:{raw.target_qualname}",
+                    path=path,
+                    line=raw.line,
+                )
+            )
+
+    # ---- 8. Entry points ---------------------------------------------
+    all_decorators = [
+        raw
+        for _path, (_syms, _imps, raw_decs) in per_file.items()
+        for raw in raw_decs
+    ]
+    entry_points = extract_entry_points(
+        repo_root=repo_root,
+        files=files,
+        decorators=all_decorators,
+        decorator_path_lookup=decorator_path_lookup,
+    )
+
+    # ---- 9. Clusters --------------------------------------------------
+    cluster_entries = _build_clusters(file_entries, py_files)
+
+    git_commit = _current_git_commit(repo_root)
+
+    return LoomIndex(
+        generated_at=datetime.now(UTC),
+        repo_root=str(repo_root.resolve()),
+        git_commit=git_commit,
+        files=file_entries,
+        symbols=symbol_entries,
+        imports=import_edges,
+        # Call graph stays empty in v1 — see _ast_walk's design note.
+        # Schema accommodates future call-graph extraction so we
+        # don't need a schema bump when we add it.
+        calls=list[CallEdge](),
+        decorators=decorator_entries,
+        entry_points=entry_points,
+        clusters=cluster_entries,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Internals
+# ---------------------------------------------------------------------------
+
+
+def _walk_all(
+    py_files: list[DiscoveredFile],
+) -> dict[
+    str,
+    tuple[
+        list[_RawSymbol],
+        list,  # _RawImport
+        list[_RawDecorator],
+    ],
+]:
+    """Parse every Python file in parallel-friendly form (sequential
+    for now; can wrap in ``anyio.to_thread.run_sync`` later if
+    profiling shows it matters)."""
+    out = {}
+    for f in py_files:
+        try:
+            text = f.abs_path.read_text(encoding="utf-8")
+        except OSError:
+            continue
+        syms, imps, decs = walk_python_file(text, f.rel_path)
+        out[f.rel_path] = (syms, imps, decs)
+    return out
+
+
+def _pagerank_edges(
+    edges: Iterable,  # ImportEdge
+    py_files: list[DiscoveredFile],
+    module_index,
+    raw_imports_by_file: dict[str, list[tuple[str, int, int]]],
+) -> list[tuple[str, str]]:
+    """Convert :class:`ImportEdge` records into ``(from_path,
+    to_path)`` pairs suitable for :func:`_graph.pagerank_file_graph`.
+
+    We can't read ``to_path`` straight off :class:`ImportEdge` (the
+    schema stores the display module name, not the resolved file).
+    So we re-resolve here using the same machinery the edges came
+    from. Cheap and keeps the schema clean.
+    """
+    from ._resolve import resolve_import
+
+    py_paths = {f.rel_path for f in py_files}
+    pairs: list[tuple[str, str]] = []
+    for from_path, items in raw_imports_by_file.items():
+        if from_path not in py_paths:
+            continue
+        for to_module, _line, level in items:
+            target = resolve_import(
+                from_file=from_path,
+                to_module=to_module,
+                level=level,
+                module_index=module_index,
+            )
+            if target is not None and target in py_paths:
+                pairs.append((from_path, target))
+    return pairs
+
+
+def _degree_maps(
+    edges: list[tuple[str, str]],
+) -> tuple[dict[str, int], dict[str, int]]:
+    """``in_degree[file]`` = number of files importing it.
+    ``out_degree[file]`` = number of files it imports.
+
+    Both are used for the SymbolEntry's ``n_callers`` / ``n_callees``
+    fields as file-level approximations until call-graph extraction
+    lands."""
+    in_deg: dict[str, int] = {}
+    out_deg: dict[str, int] = {}
+    for src, dst in edges:
+        out_deg[src] = out_deg.get(src, 0) + 1
+        in_deg[dst] = in_deg.get(dst, 0) + 1
+    return in_deg, out_deg
+
+
+def _build_clusters(
+    files: list[FileEntry], py_files: list[DiscoveredFile]
+) -> list[Cluster]:
+    """Group Python files into subsystem clusters via path prefix,
+    compute centroid symbols (highest-PageRank file → its symbols)
+    and hash buckets per cluster."""
+    py_paths = [f.rel_path for f in py_files]
+    groups = cluster_by_path_prefix(py_paths)
+
+    # Map files -> entries for hash lookup.
+    by_path = {f.path: f for f in files}
+
+    clusters: list[Cluster] = []
+    for cluster_id, paths in groups.items():
+        hashes = sorted(
+            by_path[p].sha256 for p in paths if p in by_path
+        )
+        bucket = hashlib.sha256(
+            "\n".join(hashes).encode("utf-8")
+        ).hexdigest()[:16]
+        title = cluster_id.replace("/", "/").rstrip("/") or "root"
+        clusters.append(
+            Cluster(
+                id=cluster_id.replace("/", ".") or "root",
+                title=title,
+                paths=list(paths),
+                centroid_symbols=[],  # populated in annotator pass
+                centrality=0.0,  # populated alongside centroids
+                hash_bucket=bucket,
+            )
+        )
+    return sorted(clusters, key=lambda c: c.id)
+
+
+def _current_git_commit(repo_root: Path) -> str | None:
+    """``git rev-parse HEAD`` — short-form (12 chars) commit hash.
+    Useful for diff-aware refresh: "the index was built against
+    commit X; you're on commit Y; here's what files changed"."""
+    if not is_git_repo(repo_root):
+        return None
+    import subprocess
+
+    try:
+        proc = subprocess.run(
+            ["git", "rev-parse", "--short=12", "HEAD"],
+            cwd=str(repo_root),
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=5,
+        )
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return None
+    if proc.returncode != 0:
+        return None
+    out = proc.stdout.strip()
+    return out or None

loom_code/loominit/repomap.py

@@ -0,0 +1,225 @@
+"""Deterministic, LLM-free repo map — an Aider-style ranked symbol
+overview.
+
+Why this exists: the older LOOM.md path injected BM25-ranked *LLM
+narrative* sections, which (a) drifted from the code the moment the
+agent edited it and (b) was a lossy paraphrase, not the code. This
+renderer instead consumes the **structural** index (``build_index`` —
+AST walk, no model calls) and emits the most structurally-important
+symbols — real signatures + locations — within a token budget. It is:
+
+* **deterministic** — same index in, same map out (stable tiebreaks);
+* **fresh-by-construction** — re-running ``build_index`` reflects the
+  current tree, with zero LLM cost (unlike re-annotation);
+* **cache-friendly** — the map is a stable global overview, so it
+  doesn't churn the system prompt per turn the way BM25 retrieval did.
+
+Ranking heuristic (Aider's core insight: "a symbol referenced by many
+others is more valuable context than a private helper called once"):
+``n_callers`` (how widely the symbol's file is imported — already
+computed by the extractor) plus bonuses for entry points, public
+surface, and classes (architecture lands first).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .schema import LoomIndex, SymbolEntry
+
+# Rough chars→tokens ratio for budgeting (English/code ≈ 4 chars/token).
+_CHARS_PER_TOKEN = 4
+
+
+_DEF_PREFIXES = ("class ", "def ", "async def ")
+
+
+def _score(sym: SymbolEntry, entry_point_paths: set[str]) -> float:
+    """Structural importance of a symbol. Higher = surfaced sooner."""
+    score = float(sym.n_callers) * 2.0  # import popularity (the core signal)
+    if not sym.name.startswith("_"):
+        score += 1.0  # public surface beats private helpers
+    if sym.signature.startswith("class "):
+        score += 2.0  # classes carry architecture
+    elif not sym.signature.startswith(_DEF_PREFIXES):
+        score -= 3.0  # bare module constants are noise in an overview
+    if sym.path in entry_point_paths:
+        score += 6.0  # CLI / route / main entry points orient fast
+    return score
+
+
+def build_repo_map(
+    index: LoomIndex,
+    *,
+    max_tokens: int = 1500,
+    max_per_file: int = 8,
+) -> str:
+    """Render the top symbols by structural importance, grouped by
+    file, within ``max_tokens`` (best-effort char estimate). Files are
+    ordered by their most-important symbol; within a file, symbols are
+    ordered by score then line. Deterministic."""
+    budget = max_tokens * _CHARS_PER_TOKEN
+    ep_paths = {ep.path for ep in index.entry_points if ep.path}
+
+    by_file: dict[str, list[SymbolEntry]] = {}
+    for sym in index.symbols:
+        by_file.setdefault(sym.path, []).append(sym)
+    for syms in by_file.values():
+        syms.sort(key=lambda s: (-_score(s, ep_paths), s.line))
+
+    def file_rank(path: str) -> float:
+        return max((_score(s, ep_paths) for s in by_file[path]), default=0.0)
+
+    file_order = sorted(by_file, key=lambda p: (-file_rank(p), p))
+
+    head = "# Repo map — top symbols by structural importance\n"
+    out: list[str] = [head]
+    used = len(head)
+    for path in file_order:
+        header = f"\n## {path}\n"
+        if used + len(header) > budget:
+            break
+        block: list[str] = [header]
+        blen = len(header)
+        for sym in by_file[path][:max_per_file]:
+            doc = (
+                f"  — {sym.docstring_first_line}"
+                if sym.docstring_first_line
+                else ""
+            )
+            entry = f"- `{sym.signature}` ({sym.path}:{sym.line}){doc}\n"
+            if used + blen + len(entry) > budget:
+                break
+            block.append(entry)
+            blen += len(entry)
+        # Only emit the file header if at least one symbol fit under it.
+        if blen > len(header):
+            out.extend(block)
+            used += blen
+    return "".join(out)
+
+
+def repo_map_for_root(
+    root: Path | str, *, max_tokens: int = 1500
+) -> str | None:
+    """Convenience: build the structural index for ``root`` (no LLM)
+    and render the repo map. Returns ``None`` when the repo has no
+    indexable symbols. Caller decides caching/freshness policy."""
+    from .extractor import build_index
+
+    index = build_index(Path(root))
+    if not index.symbols:
+        return None
+    return build_repo_map(index, max_tokens=max_tokens)
+
+
+# Dirs the freshness signature ignores — build_index skips them too, so
+# walking them would only add cost + false cache misses.
+_SKIP_DIRS = frozenset(
+    {
+        ".git",
+        ".loom",
+        ".loom-code",
+        "node_modules",
+        ".venv",
+        "venv",
+        "__pycache__",
+        ".mypy_cache",
+        ".ruff_cache",
+        "dist",
+        "build",
+        ".tox",
+    }
+)
+
+# Cache: root -> (signature, rendered map). The signature is cheap to
+# recompute; the expensive AST walk only re-runs when it changes.
+_REPO_MAP_CACHE: dict[str, tuple[tuple[float, int], str | None]] = {}
+
+
+def _tree_signature(root: Path) -> tuple[float, int]:
+    """A cheap freshness key: (newest .py mtime, file count). Either
+    moving means the tree changed → rebuild. Count catches add/delete
+    that don't bump the newest mtime."""
+    newest = 0.0
+    count = 0
+    for p in root.rglob("*.py"):
+        if any(part in _SKIP_DIRS for part in p.parts):
+            continue
+        try:
+            newest = max(newest, p.stat().st_mtime)
+            count += 1
+        except OSError:
+            continue
+    return (newest, count)
+
+
+def _repo_map_cached(
+    root: Path | str, max_tokens: int
+) -> tuple[str | None, bool]:
+    """Core cached build. Returns ``(map, rebuilt)`` — ``rebuilt`` is
+    True when the source tree changed and we re-walked, False on a
+    cache hit. Freshness-keyed (newest .py mtime + file count), so it
+    only re-parses when something actually moved. Mirrors the map to
+    ``<root>/.loom/repomap.md`` on every rebuild for inspection."""
+    root_p = Path(root).resolve()
+    key = str(root_p)
+    sig = _tree_signature(root_p)
+    cached = _REPO_MAP_CACHE.get(key)
+    if cached is not None and cached[0] == sig:
+        return cached[1], False
+    rendered = repo_map_for_root(root_p, max_tokens=max_tokens)
+    _REPO_MAP_CACHE[key] = (sig, rendered)
+    _write_repomap_file(root_p, rendered)
+    return rendered, True
+
+
+def repo_map_for_root_cached(
+    root: Path | str, *, max_tokens: int = 1500
+) -> str | None:
+    """Cached repo map (the "fresh-by-construction" path). Safe to call
+    every turn — re-walks only when the source tree changed. The agent
+    reads the map from its memory block; ``.loom/repomap.md`` mirrors it
+    for inspection."""
+    return _repo_map_cached(root, max_tokens)[0]
+
+
+def repo_map_meta_for_root_cached(
+    root: Path | str, *, max_tokens: int = 1500
+) -> dict[str, object]:
+    """Like :func:`repo_map_for_root_cached` but returns the map plus
+    metadata for UI surfaces: ``{map, rebuilt, symbols, chars}``.
+    ``rebuilt`` distinguishes a fresh re-index (tree changed) from a
+    cache hit, so the desktop can show "re-indexed" vs "cached"."""
+    body, rebuilt = _repo_map_cached(root, max_tokens)
+    text = body or ""
+    return {
+        "map": text,
+        "rebuilt": rebuilt,
+        "symbols": text.count("\n- `"),
+        "chars": len(text),
+    }
+
+
+def _write_repomap_file(root: Path, rendered: str | None) -> None:
+    """Mirror the injected repo map to ``<root>/.loom/repomap.md`` so a
+    human can see what the agent receives. Overwritten on every rebuild;
+    best-effort (a write failure never affects the returned map). The
+    agent reads the map from its memory block, not this file — it's
+    inspection-only. ``.loom/`` is gitignored, so this isn't committed."""
+    if not rendered:
+        return
+    try:
+        loom_dir = root / ".loom"
+        loom_dir.mkdir(parents=True, exist_ok=True)
+        note = (
+            "<!-- Auto-generated by loom-code. Rebuilt whenever the "
+            "source tree changes; this is the repo map injected into the "
+            "agent's system prompt each turn (the `loom_index` block). "
+            "Inspection-only — do not edit, it is overwritten. -->\n\n"
+        )
+        (loom_dir / "repomap.md").write_text(
+            note + rendered, encoding="utf-8"
+        )
+    except OSError:
+        pass