ontosight-codegraph 0.1.0__py3-none-any.whl

ontosight_codegraph/store.py

@@ -0,0 +1,589 @@
+"""Read CodeGraph SQLite index and build call-graph subgraphs for OntoSight."""
+
+from __future__ import annotations
+
+import re
+import sqlite3
+from collections import deque
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Iterable, List, Optional, Sequence, Set, Tuple
+
+from pydantic import BaseModel, Field
+
+EXCLUDED_NODE_KINDS = frozenset({"parameter", "import", "export"})
+CALL_EDGE_KIND = "calls"
+INIT_HINT = "Run: npx @colbymchenry/codegraph init -i"
+
+NODE_SELECT = """
+SELECT id, kind, name, qualified_name, file_path, language,
+       start_line, end_line, signature, docstring
+FROM nodes
+"""
+
+
+class CodeGraphNotFoundError(FileNotFoundError):
+    """Raised when `.codegraph/codegraph.db` is missing."""
+
+
+class CodeSymbolNode(BaseModel):
+    """A code symbol from the CodeGraph index."""
+
+    id: str
+    kind: str
+    name: str
+    qualified_name: str
+    file_path: str
+    language: str
+    start_line: int
+    end_line: int = 0
+    signature: Optional[str] = None
+    docstring: Optional[str] = None
+
+
+class CodeCallEdge(BaseModel):
+    """A call edge between two symbols."""
+
+    id: str
+    source_id: str
+    target_id: str
+    line: Optional[int] = None
+
+
+@dataclass(frozen=True)
+class SubgraphResult:
+    """Extracted call subgraph plus metadata for CLI / viewer."""
+
+    nodes: List[CodeSymbolNode]
+    edges: List[CodeCallEdge]
+    truncated: bool
+    filter_summary: str
+    languages: List[str]
+    seed_ids: List[str]
+
+
+def resolve_codegraph_db(project_path: Path) -> Path:
+    """Return path to CodeGraph database or raise with init hint."""
+    root = project_path.resolve()
+    db_path = root / ".codegraph" / "codegraph.db"
+    if not db_path.is_file():
+        raise CodeGraphNotFoundError(
+            f"CodeGraph index not found at {db_path}. {INIT_HINT}"
+        )
+    return db_path
+
+
+def _row_to_node(row: sqlite3.Row) -> CodeSymbolNode:
+    return CodeSymbolNode(
+        id=row["id"],
+        kind=row["kind"],
+        name=row["name"],
+        qualified_name=row["qualified_name"],
+        file_path=row["file_path"],
+        language=row["language"],
+        start_line=row["start_line"],
+        end_line=row["end_line"] or 0,
+        signature=row["signature"],
+        docstring=row["docstring"],
+    )
+
+
+def _normalize_path_prefix(path_filter: Optional[str]) -> Optional[str]:
+    if not path_filter:
+        return None
+    cleaned = path_filter.replace("\\", "/").strip("/")
+    return cleaned or None
+
+
+def _tokenize_task(task: str) -> List[str]:
+    tokens = re.findall(r"[A-Za-z_][A-Za-z0-9_]*|[A-Za-z]{3,}", task.lower())
+    return [t for t in tokens if len(t) >= 2]
+
+
+class CodeGraphStore:
+    """Read-only access to a CodeGraph SQLite database."""
+
+    def __init__(self, db_path: Path):
+        self.db_path = db_path
+        self._conn = sqlite3.connect(str(db_path))
+        self._conn.row_factory = sqlite3.Row
+
+    def close(self) -> None:
+        self._conn.close()
+
+    def __enter__(self) -> "CodeGraphStore":
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
+    def _path_clause(self, path_filter: Optional[str]) -> Tuple[str, List[str]]:
+        prefix = _normalize_path_prefix(path_filter)
+        if not prefix:
+            return "", []
+        return " AND file_path LIKE ? ESCAPE '\\'", [f"{prefix}%"]
+
+    def _eligible_node_clause(self) -> str:
+        placeholders = ",".join("?" * len(EXCLUDED_NODE_KINDS))
+        return f" AND kind NOT IN ({placeholders})"
+
+    def fetch_nodes_by_ids(self, node_ids: Set[str]) -> List[CodeSymbolNode]:
+        if not node_ids:
+            return []
+        placeholders = ",".join("?" * len(node_ids))
+        query = f"{NODE_SELECT} WHERE id IN ({placeholders})"
+        rows = self._conn.execute(query, list(node_ids)).fetchall()
+        return [_row_to_node(r) for r in rows]
+
+    def fetch_call_edges_for_nodes(self, node_ids: Set[str]) -> List[CodeCallEdge]:
+        if not node_ids:
+            return []
+        placeholders = ",".join("?" * len(node_ids))
+        query = f"""
+            SELECT rowid AS edge_rowid, source, target, line
+            FROM edges
+            WHERE kind = ?
+              AND source IN ({placeholders})
+              AND target IN ({placeholders})
+        """
+        params: List[object] = [CALL_EDGE_KIND, *node_ids, *node_ids]
+        rows = self._conn.execute(query, params).fetchall()
+        edges: List[CodeCallEdge] = []
+        for row in rows:
+            edges.append(
+                CodeCallEdge(
+                    id=f"call:{row['source']}:{row['target']}:{row['edge_rowid']}",
+                    source_id=row["source"],
+                    target_id=row["target"],
+                    line=row["line"],
+                )
+            )
+        return edges
+
+    def _load_adjacency(
+        self, path_filter: Optional[str]
+    ) -> Tuple[dict[str, Set[str]], dict[str, int]]:
+        """Build undirected adjacency and in-degree for call edges within path filter."""
+        path_sql, path_params = self._path_clause(path_filter)
+        kind_params = list(EXCLUDED_NODE_KINDS)
+        node_query = f"""
+            SELECT id FROM nodes
+            WHERE 1=1{self._eligible_node_clause()}{path_sql}
+        """
+        allowed = {
+            row["id"]
+            for row in self._conn.execute(node_query, kind_params + path_params).fetchall()
+        }
+        if not allowed:
+            return {}, {}
+
+        placeholders = ",".join("?" * len(allowed))
+        edge_query = f"""
+            SELECT source, target FROM edges
+            WHERE kind = ?
+              AND source IN ({placeholders})
+              AND target IN ({placeholders})
+        """
+        params: List[object] = [CALL_EDGE_KIND, *allowed, *allowed]
+        adjacency: dict[str, Set[str]] = {nid: set() for nid in allowed}
+        in_degree: dict[str, int] = {nid: 0 for nid in allowed}
+
+        for row in self._conn.execute(edge_query, params).fetchall():
+            src, tgt = row["source"], row["target"]
+            adjacency[src].add(tgt)
+            adjacency[tgt].add(src)
+            in_degree[tgt] = in_degree.get(tgt, 0) + 1
+
+        return adjacency, in_degree
+
+    def find_symbol_seeds(self, symbol: str, path_filter: Optional[str]) -> List[str]:
+        path_sql, path_params = self._path_clause(path_filter)
+        kind_params = list(EXCLUDED_NODE_KINDS)
+        exact_query = f"""
+            SELECT id FROM nodes
+            WHERE (name = ? OR qualified_name = ?)
+            {self._eligible_node_clause()}{path_sql}
+            ORDER BY start_line
+            LIMIT 10
+        """
+        params: List[object] = [symbol, symbol, *kind_params, *path_params]
+        rows = self._conn.execute(exact_query, params).fetchall()
+        if rows:
+            return [row["id"] for row in rows]
+
+        like_query = f"""
+            SELECT id FROM nodes
+            WHERE (name LIKE ? OR qualified_name LIKE ?)
+            {self._eligible_node_clause()}{path_sql}
+            ORDER BY start_line
+            LIMIT 10
+        """
+        like = f"%{symbol}%"
+        params = [like, like, *kind_params, *path_params]
+        rows = self._conn.execute(like_query, params).fetchall()
+        return [row["id"] for row in rows]
+
+    def find_task_seeds(self, task: str, path_filter: Optional[str], limit: int = 5) -> List[str]:
+        tokens = _tokenize_task(task)
+        if not tokens:
+            return self.find_symbol_seeds(task, path_filter)
+
+        scores: dict[str, int] = {}
+        path_sql, path_params = self._path_clause(path_filter)
+        kind_params = list(EXCLUDED_NODE_KINDS)
+
+        for token in tokens:
+            query = f"""
+                SELECT id, name, qualified_name, docstring FROM nodes
+                WHERE (LOWER(name) LIKE ? OR LOWER(qualified_name) LIKE ?
+                       OR LOWER(COALESCE(docstring, '')) LIKE ?)
+                {self._eligible_node_clause()}{path_sql}
+                LIMIT 50
+            """
+            pattern = f"%{token}%"
+            params: List[object] = [pattern, pattern, pattern, *kind_params, *path_params]
+            for row in self._conn.execute(query, params).fetchall():
+                nid = row["id"]
+                text = " ".join(
+                    filter(
+                        None,
+                        [row["name"], row["qualified_name"], row["docstring"] or ""],
+                    )
+                ).lower()
+                score = sum(1 for t in tokens if t in text)
+                scores[nid] = scores.get(nid, 0) + score
+
+        ranked = sorted(scores.items(), key=lambda item: (-item[1], item[0]))
+        return [nid for nid, _ in ranked[:limit]]
+
+    def pick_auto_seeds(self, path_filter: Optional[str], limit: int = 3) -> List[str]:
+        _, in_degree = self._load_adjacency(path_filter)
+        if not in_degree:
+            return []
+        ranked = sorted(in_degree.items(), key=lambda item: (-item[1], item[0]))
+        return [nid for nid, deg in ranked[:limit] if deg > 0] or [ranked[0][0]]
+
+    def search_symbols(
+        self, query: str, path_filter: Optional[str], limit: int = 10
+    ) -> List[CodeSymbolNode]:
+        path_sql, path_params = self._path_clause(path_filter)
+        kind_params = list(EXCLUDED_NODE_KINDS)
+        pattern = f"%{query.strip()}%"
+        sql = f"""
+            SELECT id, kind, name, qualified_name, file_path, language,
+                   start_line, end_line, signature, docstring
+            FROM nodes
+            WHERE (name LIKE ? OR qualified_name LIKE ?
+                   OR COALESCE(signature, '') LIKE ?
+                   OR COALESCE(docstring, '') LIKE ?)
+            {self._eligible_node_clause()}{path_sql}
+            ORDER BY
+                CASE WHEN name = ? THEN 0
+                     WHEN name LIKE ? THEN 1
+                     ELSE 2 END,
+                start_line
+            LIMIT ?
+        """
+        exact = query.strip()
+        params: List[object] = [
+            pattern,
+            pattern,
+            pattern,
+            pattern,
+            *kind_params,
+            *path_params,
+            exact,
+            f"{exact}%",
+            limit,
+        ]
+        rows = self._conn.execute(sql, params).fetchall()
+        return [_row_to_node(r) for r in rows]
+
+    def expand_subgraph(
+        self,
+        seed_ids: Sequence[str],
+        *,
+        path_filter: Optional[str],
+        hops: int,
+        max_nodes: int,
+    ) -> Tuple[Set[str], bool]:
+        adjacency, _ = self._load_adjacency(path_filter)
+        if not adjacency:
+            return set(seed_ids), False
+
+        visited: Set[str] = set()
+        queue: deque[Tuple[str, int]] = deque()
+
+        for seed in seed_ids:
+            if seed in adjacency:
+                visited.add(seed)
+                queue.append((seed, 0))
+
+        if not visited:
+            # Seeds may be outside path filter; include them alone.
+            return set(seed_ids), False
+
+        truncated = False
+        while queue:
+            node_id, depth = queue.popleft()
+            if depth >= hops:
+                continue
+            for neighbor in adjacency.get(node_id, ()):
+                if neighbor in visited:
+                    continue
+                if len(visited) >= max_nodes:
+                    truncated = True
+                    return visited, truncated
+                visited.add(neighbor)
+                queue.append((neighbor, depth + 1))
+
+        return visited, truncated
+
+
+def _build_filter_summary(
+    *,
+    path_filter: Optional[str],
+    symbol: Optional[str],
+    task: Optional[str],
+    seed_ids: Sequence[str],
+) -> str:
+    parts: List[str] = []
+    if path_filter:
+        parts.append(f"path={path_filter}")
+    if symbol:
+        parts.append(f"symbol={symbol}")
+    elif task:
+        parts.append(f"task={task!r}")
+    elif seed_ids:
+        parts.append(f"seeds={len(seed_ids)}")
+    else:
+        parts.append("auto-seed")
+    return ", ".join(parts)
+
+
+def load_call_subgraph(
+    project_path: Path,
+    *,
+    path_filter: Optional[str] = None,
+    symbol: Optional[str] = None,
+    task: Optional[str] = None,
+    hops: int = 2,
+    max_nodes: int = 200,
+) -> SubgraphResult:
+    """Load a bounded call-graph subgraph from CodeGraph."""
+    if hops < 0:
+        raise ValueError("hops must be >= 0")
+    if max_nodes < 1:
+        raise ValueError("max_nodes must be >= 1")
+    if symbol and task:
+        raise ValueError("Specify only one of symbol or task")
+
+    db_path = resolve_codegraph_db(project_path)
+
+    with CodeGraphStore(db_path) as store:
+        if symbol:
+            seed_ids = store.find_symbol_seeds(symbol, path_filter)
+            if not seed_ids:
+                raise ValueError(f"No symbol matching {symbol!r} in CodeGraph index")
+        elif task:
+            seed_ids = store.find_task_seeds(task, path_filter)
+            if not seed_ids:
+                raise ValueError(f"No symbols matching task {task!r} in CodeGraph index")
+        else:
+            seed_ids = store.pick_auto_seeds(path_filter)
+
+        if not seed_ids:
+            return SubgraphResult(
+                nodes=[],
+                edges=[],
+                truncated=False,
+                filter_summary=_build_filter_summary(
+                    path_filter=path_filter,
+                    symbol=symbol,
+                    task=task,
+                    seed_ids=seed_ids,
+                ),
+                languages=[],
+                seed_ids=[],
+            )
+
+        node_ids, truncated = store.expand_subgraph(
+            seed_ids,
+            path_filter=path_filter,
+            hops=hops,
+            max_nodes=max_nodes,
+        )
+        # Always include seeds even if path filter excluded them from adjacency.
+        node_ids.update(seed_ids)
+
+        nodes = store.fetch_nodes_by_ids(node_ids)
+        # Drop excluded kinds unless they are call endpoints.
+        edges = store.fetch_call_edges_for_nodes(node_ids)
+        endpoint_ids = {e.source_id for e in edges} | {e.target_id for e in edges}
+        nodes = [
+            n
+            for n in nodes
+            if n.kind not in EXCLUDED_NODE_KINDS or n.id in endpoint_ids
+        ]
+        allowed_ids = {n.id for n in nodes}
+        edges = [
+            e
+            for e in edges
+            if e.source_id in allowed_ids and e.target_id in allowed_ids
+        ]
+        languages = sorted({n.language for n in nodes if n.language})
+
+        return SubgraphResult(
+            nodes=nodes,
+            edges=edges,
+            truncated=truncated,
+            filter_summary=_build_filter_summary(
+                path_filter=path_filter,
+                symbol=symbol,
+                task=task,
+                seed_ids=seed_ids,
+            ),
+            languages=languages,
+            seed_ids=list(seed_ids),
+        )
+
+
+_QUERY_KEY_ALIASES = {
+    "path": "path_filter",
+    "max-nodes": "max_nodes",
+    "max_nodes": "max_nodes",
+}
+
+
+def parse_codegraph_query(query: str) -> dict:
+    """Parse CLI-style CodeGraph query strings into load_call_subgraph kwargs.
+
+    Supported forms:
+      symbol=load_call_subgraph path=hyperextract/ hops=3 max_nodes=500
+      symbol:load_call_subgraph task="fix parser bug"
+    """
+    if not query or not query.strip():
+        return {}
+
+    text = query.strip()
+    result: dict = {}
+    i = 0
+    length = len(text)
+
+    def skip_ws() -> None:
+        nonlocal i
+        while i < length and text[i].isspace():
+            i += 1
+
+    def read_quoted() -> str:
+        nonlocal i
+        quote = text[i]
+        i += 1
+        start = i
+        while i < length and text[i] != quote:
+            if text[i] == "\\" and i + 1 < length:
+                i += 2
+                continue
+            i += 1
+        value = text[start:i]
+        if i < length:
+            i += 1
+        return value
+
+    def read_token() -> str:
+        nonlocal i
+        start = i
+        while i < length and not text[i].isspace() and text[i] not in "=:":
+            i += 1
+        return text[start:i]
+
+    while i < length:
+        skip_ws()
+        if i >= length:
+            break
+
+        key = read_token()
+        if not key:
+            i += 1
+            continue
+
+        skip_ws()
+        if i >= length or text[i] not in "=:":
+            raise ValueError(f"Expected '=' or ':' after key {key!r}")
+        i += 1
+        skip_ws()
+
+        if i < length and text[i] in "\"'":
+            value = read_quoted()
+        else:
+            value = read_token()
+
+        canonical = _QUERY_KEY_ALIASES.get(key.lower(), key.lower())
+        if canonical in ("hops", "max_nodes"):
+            try:
+                result[canonical] = int(value)
+            except ValueError as exc:
+                raise ValueError(f"Invalid integer for {key}: {value!r}") from exc
+        else:
+            result[canonical] = value
+
+    if "symbol" in result and "task" in result:
+        raise ValueError("Specify only one of symbol or task")
+    return result
+
+
+def merge_query_params(
+    *,
+    query: Optional[str] = None,
+    path_filter: Optional[str] = None,
+    symbol: Optional[str] = None,
+    task: Optional[str] = None,
+    hops: Optional[int] = None,
+    max_nodes: Optional[int] = None,
+    default_path_filter: Optional[str] = None,
+) -> dict:
+    """Merge structured fields with an optional query string (structured wins)."""
+    merged: dict = {}
+    if query and query.strip():
+        merged.update(parse_codegraph_query(query))
+
+    if path_filter is not None:
+        merged["path_filter"] = path_filter
+    elif default_path_filter and "path_filter" not in merged:
+        merged["path_filter"] = default_path_filter
+    if symbol is not None:
+        merged["symbol"] = symbol
+    if task is not None:
+        merged["task"] = task
+    if hops is not None:
+        merged["hops"] = hops
+    if max_nodes is not None:
+        merged["max_nodes"] = max_nodes
+
+    if merged.get("symbol") and merged.get("task"):
+        raise ValueError("Specify only one of symbol or task")
+    return merged
+
+
+def make_search_callback(
+    project_path: Path,
+    path_filter: Optional[str] = None,
+    *,
+    top_k: int = 10,
+) -> Callable[[str], Tuple[List[CodeSymbolNode], List[CodeCallEdge]]]:
+    """Create OntoSight search callback backed by CodeGraph symbol search."""
+
+    def search(query: str) -> Tuple[List[CodeSymbolNode], List[CodeCallEdge]]:
+        db_path = resolve_codegraph_db(project_path)
+        with CodeGraphStore(db_path) as store:
+            nodes = store.search_symbols(query, path_filter, limit=top_k)
+            node_ids = {n.id for n in nodes}
+            edges = store.fetch_call_edges_for_nodes(node_ids)
+            edges = [
+                e
+                for e in edges
+                if e.source_id in node_ids and e.target_id in node_ids
+            ]
+            return nodes, edges
+
+    return search