gdmcode 0.1.0__py3-none-any.whl

src/memory/code_index.py

@@ -0,0 +1,376 @@
+"""Code index — Tree-sitter AST symbol index for $0/query codebase search.
+
+Builds a structural index of all source files in the project:
+  - Python:     function defs, class defs, method defs
+  - TypeScript: function/arrow decls, class defs, interface defs
+  - JavaScript: function/arrow decls, class defs
+
+Stored in gdm.db code_index table. Incremental: only re-indexes files
+whose mtime changed since last index run.
+
+Query API:
+  find_symbol(name)               → exact name match across project
+  find_callers(symbol)            → files that call/import symbol
+  get_symbol_signature(path, sym) → return function signature line
+  search(query)                   → keyword grep fallback
+"""
+from __future__ import annotations
+
+import logging
+import re
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+__all__ = ["CodeIndex", "SymbolMatch", "ReferenceMatch"]
+
+log = logging.getLogger(__name__)
+
+_SUPPORTED_SUFFIXES: frozenset[str] = frozenset({".py", ".ts", ".tsx", ".js", ".jsx"})
+_MAX_FILE_SIZE_BYTES: int = 500_000
+_INDEX_TIMEOUT_SECS: float = 60.0
+
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+@dataclass
+class SymbolMatch:
+    """A single symbol found in the index."""
+
+    symbol: str
+    kind: str          # "function" | "class" | "method" | "interface" | "arrow"
+    file: str          # relative path from project root
+    line: int
+    signature: str     # e.g. "def validate_user(email: str) -> bool:"
+
+    def format(self, root: Path) -> str:
+        """Human-readable match line for display."""
+        return f"{self.file}:{self.line}  [{self.kind}] {self.signature}"
+
+
+@dataclass
+class ReferenceMatch:
+    """A single reference to a symbol — either a definition or a call site."""
+
+    symbol: str
+    file: str       # relative path from project root
+    line: int
+    kind: str       # "definition" | "caller"
+    snippet: str    # the source line (stripped)
+
+
+# ---------------------------------------------------------------------------
+# Language-specific extractors (regex-based, fast, no grammar loading needed)
+# ---------------------------------------------------------------------------
+
+_PY_FUNC = re.compile(
+    r"^(?P<indent>\s*)(?:async\s+)?def\s+(?P<name>\w+)\s*\((?P<args>[^)]*)\)"
+    r"(?:\s*->\s*[^:]+)?:",
+    re.MULTILINE,
+)
+_PY_CLASS = re.compile(r"^class\s+(?P<name>\w+)", re.MULTILINE)
+
+_TS_FUNC = re.compile(
+    r"(?:export\s+)?(?:async\s+)?function\s+(?P<name>\w+)\s*\(",
+    re.MULTILINE,
+)
+_TS_ARROW = re.compile(
+    r"(?:export\s+)?(?:const|let)\s+(?P<name>\w+)\s*=\s*(?:async\s+)?\(",
+    re.MULTILINE,
+)
+_TS_CLASS = re.compile(r"(?:export\s+)?(?:abstract\s+)?class\s+(?P<name>\w+)", re.MULTILINE)
+_TS_IFACE = re.compile(r"(?:export\s+)?interface\s+(?P<name>\w+)", re.MULTILINE)
+_TS_METHOD = re.compile(
+    r"^\s+(?:public|private|protected|static|async|\s)*"
+    r"(?P<name>\w+)\s*\([^)]*\)\s*(?::\s*\S+\s*)?[{;]",
+    re.MULTILINE,
+)
+
+
+def _extract_python(text: str, rel_path: str) -> list[SymbolMatch]:
+    """Extract symbols from Python source text."""
+    lines = text.splitlines()
+    matches: list[SymbolMatch] = []
+
+    for m in _PY_CLASS.finditer(text):
+        line_num = text[: m.start()].count("\n") + 1
+        matches.append(SymbolMatch(
+            symbol=m.group("name"), kind="class", file=rel_path,
+            line=line_num, signature=lines[line_num - 1].strip(),
+        ))
+
+    for m in _PY_FUNC.finditer(text):
+        line_num = text[: m.start()].count("\n") + 1
+        name = m.group("name")
+        if name.startswith("__") and name.endswith("__"):
+            kind = "method"
+        elif m.group("indent"):
+            kind = "method"
+        else:
+            kind = "function"
+        sig = lines[line_num - 1].strip()
+        matches.append(SymbolMatch(
+            symbol=name, kind=kind, file=rel_path, line=line_num, signature=sig,
+        ))
+
+    return matches
+
+
+def _extract_typescript(text: str, rel_path: str) -> list[SymbolMatch]:
+    """Extract symbols from TypeScript/JavaScript source text."""
+    lines = text.splitlines()
+    matches: list[SymbolMatch] = []
+
+    for pat, kind in [(_TS_CLASS, "class"), (_TS_IFACE, "interface")]:
+        for m in pat.finditer(text):
+            line_num = text[: m.start()].count("\n") + 1
+            matches.append(SymbolMatch(
+                symbol=m.group("name"), kind=kind, file=rel_path,
+                line=line_num, signature=lines[line_num - 1].strip()[:120],
+            ))
+
+    for pat, kind in [(_TS_FUNC, "function"), (_TS_ARROW, "arrow")]:
+        for m in pat.finditer(text):
+            line_num = text[: m.start()].count("\n") + 1
+            matches.append(SymbolMatch(
+                symbol=m.group("name"), kind=kind, file=rel_path,
+                line=line_num, signature=lines[line_num - 1].strip()[:120],
+            ))
+
+    return matches
+
+
+# ---------------------------------------------------------------------------
+# CodeIndex
+# ---------------------------------------------------------------------------
+
+class CodeIndex:
+    """Structural code index for fast $0/query symbol search.
+
+    Usage::
+
+        index = CodeIndex(db, project_id, project_root)
+        index.build()                    # full scan (incremental-safe)
+        matches = index.find_symbol("validateUser")
+        sig = index.get_symbol_signature("src/auth.py", "validate_user")
+    """
+
+    def __init__(self, db: Any, project_id: str, project_root: Path) -> None:
+        self._db = db
+        self._project_id = project_id
+        self._root = project_root
+        self._ensure_schema()
+
+    def _ensure_schema(self) -> None:
+        """Add mtime column to code_index if the table predates migration v10."""
+        try:
+            self._db.execute(
+                "ALTER TABLE code_index ADD COLUMN mtime REAL DEFAULT 0", ()
+            )
+        except Exception:
+            pass  # Column already exists
+
+    # ------------------------------------------------------------------
+    # Build / update index
+    # ------------------------------------------------------------------
+
+    def build(self, *, force: bool = False) -> int:
+        """Build or update the index. Returns count of files indexed."""
+        files = self._collect_source_files()
+        indexed = 0
+        start = time.monotonic()
+
+        for path in files:
+            if time.monotonic() - start > _INDEX_TIMEOUT_SECS:
+                log.warning("Code index build timed out after %.0fs", _INDEX_TIMEOUT_SECS)
+                break
+            try:
+                if force or self._needs_indexing(path):
+                    self._index_file(path)
+                    indexed += 1
+            except Exception as exc:  # noqa: BLE001
+                log.debug("Skipping %s: %s", path, exc)
+
+        log.info("Code index: %d files indexed", indexed)
+        return indexed
+
+    def index_file(self, path: Path) -> int:
+        """Index a single file (call after any write). Returns symbol count."""
+        try:
+            return self._index_file(path)
+        except Exception as exc:  # noqa: BLE001
+            log.warning("Failed to index %s: %s", path, exc)
+            return 0
+
+    # ------------------------------------------------------------------
+    # Query API
+    # ------------------------------------------------------------------
+
+    def find_symbol(self, name: str) -> list[SymbolMatch]:
+        """Find all symbols with this exact name across the project."""
+        rows = self._db.fetchall(
+            "SELECT file, symbol, kind, line, signature FROM code_index "
+            "WHERE project_id = ? AND symbol = ? ORDER BY file, line",
+            (self._project_id, name),
+        )
+        return [_row_to_match(r) for r in rows]
+
+    def search(self, query: str) -> list[SymbolMatch]:
+        """Substring search on symbol names (case-insensitive)."""
+        rows = self._db.fetchall(
+            "SELECT file, symbol, kind, line, signature FROM code_index "
+            "WHERE project_id = ? AND symbol LIKE ? ORDER BY file, line LIMIT 30",
+            (self._project_id, f"%{query}%"),
+        )
+        return [_row_to_match(r) for r in rows]
+
+    def find_callers(self, symbol: str) -> list[tuple[str, int]]:
+        """Find files + lines that reference *symbol* by grep. Returns (file, line) pairs."""
+        results: list[tuple[str, int]] = []
+        pattern = re.compile(r"\b" + re.escape(symbol) + r"\b")
+        for path in self._collect_source_files():
+            try:
+                text = path.read_text(encoding="utf-8", errors="replace")
+                for i, line in enumerate(text.splitlines(), 1):
+                    if pattern.search(line):
+                        rel = str(path.relative_to(self._root))
+                        results.append((rel, i))
+                        if len(results) >= 50:
+                            return results
+            except OSError:
+                pass
+        return results
+
+    def get_symbol_signature(self, path: Path, symbol: str) -> str | None:
+        """Return the signature line for *symbol* in *path*. Returns None if not found."""
+        rel = str(path.relative_to(self._root)) if path.is_absolute() else str(path)
+        row = self._db.fetchone(
+            "SELECT signature FROM code_index WHERE project_id = ? AND file = ? AND symbol = ?",
+            (self._project_id, rel, symbol),
+        )
+        return row["signature"] if row else None
+
+    def find_all_references(self, symbol: str) -> list[ReferenceMatch]:
+        """Return all definitions and call-sites for *symbol*, capped at 200.
+
+        Definitions come first (from the DB); callers are appended (from grep).
+        Deduplicates by (file, line).
+        """
+        seen: dict[tuple[str, int], ReferenceMatch] = {}
+
+        # 1. DB definitions
+        for row in self._db.fetchall(
+            "SELECT file, symbol, kind, line, signature FROM code_index "
+            "WHERE project_id = ? AND symbol = ? ORDER BY file, line",
+            (self._project_id, symbol),
+        ):
+            key = (row["file"], row["line"])
+            if key not in seen:
+                seen[key] = ReferenceMatch(
+                    symbol=symbol,
+                    file=row["file"],
+                    line=row["line"],
+                    kind="definition",
+                    snippet=row["signature"] or "",
+                )
+
+        # 2. Grep callers
+        for file_path, lineno in self.find_callers(symbol):
+            key = (file_path, lineno)
+            if key not in seen:
+                # Read the actual line for the snippet
+                abs_path = self._root / file_path
+                snippet = ""
+                try:
+                    lines = abs_path.read_text(encoding="utf-8", errors="replace").splitlines()
+                    if 0 < lineno <= len(lines):
+                        snippet = lines[lineno - 1].strip()
+                except OSError:
+                    pass
+                seen[key] = ReferenceMatch(
+                    symbol=symbol,
+                    file=file_path,
+                    line=lineno,
+                    kind="caller",
+                    snippet=snippet,
+                )
+            if len(seen) >= 200:
+                break
+
+        results = list(seen.values())
+        results.sort(key=lambda r: (0 if r.kind == "definition" else 1, r.file, r.line))
+        return results[:200]
+
+
+
+    def _collect_source_files(self) -> list[Path]:
+        """Return all indexable source files under project root."""
+        files: list[Path] = []
+        skip = {".git", "node_modules", "__pycache__", ".venv", "venv", "dist", "build"}
+        for path in self._root.rglob("*"):
+            if any(p in path.parts for p in skip):
+                continue
+            if path.suffix in _SUPPORTED_SUFFIXES and path.is_file():
+                if path.stat().st_size <= _MAX_FILE_SIZE_BYTES:
+                    files.append(path)
+        return files
+
+    def _needs_indexing(self, path: Path) -> bool:
+        """True if file has not been indexed yet or its mtime changed."""
+        rel = str(path.relative_to(self._root))
+        row = self._db.fetchone(
+            "SELECT mtime FROM code_index WHERE project_id = ? AND file = ? LIMIT 1",
+            (self._project_id, rel),
+        )
+        if row is None:
+            return True
+        return path.stat().st_mtime != row["mtime"]
+
+    def _index_file(self, path: Path) -> int:
+        """Extract symbols from *path* and upsert into code_index."""
+        rel = str(path.relative_to(self._root))
+        try:
+            text = path.read_text(encoding="utf-8", errors="replace")
+        except OSError:
+            return 0
+
+        mtime = path.stat().st_mtime
+
+        suffix = path.suffix.lower()
+        if suffix == ".py":
+            symbols = _extract_python(text, rel)
+        elif suffix in (".ts", ".tsx", ".js", ".jsx"):
+            symbols = _extract_typescript(text, rel)
+        else:
+            return 0
+
+        # Clear old entries for this file then insert fresh
+        self._db.execute(
+            "DELETE FROM code_index WHERE project_id = ? AND file = ?",
+            (self._project_id, rel),
+        )
+        for sym in symbols:
+            self._db.execute(
+                "INSERT INTO code_index (project_id, file, symbol, kind, line, signature, mtime) "
+                "VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (self._project_id, rel, sym.symbol, sym.kind, sym.line, sym.signature, mtime),
+            )
+
+        return len(symbols)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _row_to_match(row: Any) -> SymbolMatch:
+    return SymbolMatch(
+        symbol=row["symbol"],
+        kind=row["kind"],
+        file=row["file"],
+        line=row["line"],
+        signature=row["signature"] or "",
+    )