codegraph-py 1.0.0__py3-none-any.whl

codegraph/context/__init__.py

@@ -0,0 +1,216 @@
+"""
+CodeGraph Context Builder
+
+Builds rich context for AI agents to understand code.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Set, Any, Tuple
+
+from codegraph.types import (
+    Node, Edge, Subgraph, Context, CodeBlock, TaskContext,
+    BuildContextOptions, SearchOptions,
+)
+from codegraph.db.queries import QueryBuilder
+from codegraph.graph import GraphTraverser
+
+
+class ContextBuilder:
+    """Builds rich code context for AI agents."""
+
+    def __init__(self, project_root: str, queries: QueryBuilder,
+                 traverser: GraphTraverser):
+        self._project_root = project_root
+        self._queries = queries
+        self._traverser = traverser
+
+    def build_context(self, node_id: str, depth: int = 2) -> Context:
+        """Build full context around a focal node."""
+        focal = self._queries.get_node_by_id(node_id)
+        if not focal:
+            return Context()
+
+        ctx = Context(focal=focal)
+
+        # Get ancestors
+        ancestors = self._traverser.get_ancestors(node_id)
+        ctx.ancestors = ancestors
+
+        # Get children
+        children = self._traverser.get_children(node_id)
+        ctx.children = children
+
+        # Get incoming references
+        incoming_edges = self._queries.get_incoming_edges(node_id)
+        ctx.incoming_refs = []
+        for edge in incoming_edges:
+            source = self._queries.get_node_by_id(edge.source)
+            if source:
+                ctx.incoming_refs.append({'node': source, 'edge': edge})
+
+        # Get outgoing references
+        outgoing_edges = self._queries.get_outgoing_edges(node_id)
+        ctx.outgoing_refs = []
+        for edge in outgoing_edges:
+            target = self._queries.get_node_by_id(edge.target)
+            if target:
+                ctx.outgoing_refs.append({'node': target, 'edge': edge})
+
+        # Collect types (extends/implements targets)
+        for edge in outgoing_edges:
+            if edge.kind in ('extends', 'implements', 'type_of', 'returns'):
+                target = self._queries.get_node_by_id(edge.target)
+                if target and target not in ctx.types:
+                    ctx.types.append(target)
+
+        # Collect imports
+        for edge in outgoing_edges:
+            if edge.kind == 'imports':
+                target = self._queries.get_node_by_id(edge.target)
+                if target:
+                    ctx.imports.append(target)
+
+        return ctx
+
+    def build_task_context(self, query: str,
+                           options: Optional[BuildContextOptions] = None) -> TaskContext:
+        """Build context for a task query."""
+        opts = options or BuildContextOptions()
+        task_ctx = TaskContext(query=query)
+
+        # Search for relevant nodes
+        results = self._queries.search_nodes(query, SearchOptions(limit=opts.search_limit))
+        entry_points = [r.node for r in results]
+        task_ctx.entry_points = entry_points
+
+        if not entry_points:
+            return task_ctx
+
+        # Build subgraph by traversing from entry points
+        subgraph = Subgraph()
+        visited_nodes: Set[str] = set()
+        all_edges: List[Edge] = []
+
+        for ep in entry_points:
+            if ep.id not in visited_nodes:
+                subgraph.nodes[ep.id] = ep
+                visited_nodes.add(ep.id)
+
+                # Traverse outgoing edges
+                if opts.traversal_depth > 0:
+                    edges = self._queries.get_outgoing_edges(ep.id)
+                    for edge in edges:
+                        all_edges.append(edge)
+                        if edge.target not in visited_nodes:
+                            target = self._queries.get_node_by_id(edge.target)
+                            if target:
+                                subgraph.nodes[edge.target] = target
+                                visited_nodes.add(edge.target)
+
+                # Traverse incoming edges
+                edges_in = self._queries.get_incoming_edges(ep.id)
+                for edge in edges_in:
+                    all_edges.append(edge)
+
+        # Deduplicate edges
+        seen_edges: Set[Tuple[str, str, str]] = set()
+        for edge in all_edges:
+            key = (edge.source, edge.target, edge.kind)
+            if key not in seen_edges:
+                seen_edges.add(key)
+                subgraph.edges.append(edge)
+
+        subgraph.roots = [ep.id for ep in entry_points]
+        subgraph.confidence = 'high' if entry_points else 'low'
+        task_ctx.subgraph = subgraph
+
+        # Generate code blocks from nodes
+        if opts.include_code:
+            code_blocks = []
+            nodes_to_include = list(subgraph.nodes.values())[:opts.max_code_blocks]
+            for node in nodes_to_include:
+                block = self._read_code_block(node, opts.max_code_block_size)
+                if block:
+                    code_blocks.append(block)
+            task_ctx.code_blocks = code_blocks
+
+        # Collect related files
+        files: Set[str] = set()
+        for node in subgraph.nodes.values():
+            files.add(node.file_path)
+        task_ctx.related_files = list(files)[:50]
+
+        # Generate summary
+        task_ctx.summary = self._generate_summary(query, entry_points, subgraph)
+
+        # Stats
+        task_ctx.stats = {
+            'nodeCount': len(subgraph.nodes),
+            'edgeCount': len(subgraph.edges),
+            'fileCount': len(task_ctx.related_files),
+            'codeBlockCount': len(task_ctx.code_blocks),
+            'totalCodeSize': sum(len(b.content) for b in task_ctx.code_blocks),
+        }
+
+        return task_ctx
+
+    def _read_code_block(self, node: Node, max_size: int) -> Optional[CodeBlock]:
+        """Read the source code for a node."""
+        try:
+            filepath = os.path.join(self._project_root, node.file_path)
+            if not os.path.isfile(filepath):
+                return None
+
+            with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
+                lines = f.readlines()
+
+            # Extract the range
+            start = max(0, node.start_line - 1)
+            end = min(len(lines), node.end_line)
+            code_lines = lines[start:end]
+
+            content = ''.join(code_lines).strip()
+            if not content:
+                return None
+
+            # Truncate if too long
+            if len(content) > max_size:
+                content = content[:max_size] + '\n... (truncated)'
+
+            return CodeBlock(
+                content=content,
+                file_path=node.file_path,
+                start_line=node.start_line,
+                end_line=node.end_line,
+                language=node.language,
+                node=node,
+            )
+        except Exception:
+            return None
+
+    def _generate_summary(self, query: str, entry_points: List[Node],
+                           subgraph: Subgraph) -> str:
+        """Generate a text summary of the context."""
+        if not entry_points:
+            return f'No relevant symbols found for: {query}'
+
+        ep_names = ', '.join(
+            f'{n.name} ({n.kind})' for n in entry_points[:5]
+        )
+
+        files = set(n.file_path for n in subgraph.nodes.values())
+
+        return (
+            f'Found {len(subgraph.nodes)} symbols across {len(files)} files '
+            f'related to "{query}". '
+            f'Entry points: {ep_names}.'
+        )
+
+
+import os
+
+def create_context_builder(project_root: str, queries: QueryBuilder,
+                            traverser: GraphTraverser) -> ContextBuilder:
+    """Factory function to create a ContextBuilder."""
+    return ContextBuilder(project_root, queries, traverser)

codegraph/db/__init__.py

@@ -0,0 +1,11 @@
+"""CodeGraph Database Layer."""
+
+from codegraph.db.connection import DatabaseConnection, get_database_path, remove_database_files
+from codegraph.db.queries import QueryBuilder
+
+__all__ = [
+    'DatabaseConnection',
+    'get_database_path',
+    'remove_database_files',
+    'QueryBuilder',
+]

codegraph/db/connection.py

@@ -0,0 +1,163 @@
+"""
+CodeGraph Database Connection
+
+Manages SQLite database connections and lifecycle.
+"""
+
+from __future__ import annotations
+
+import os
+import sqlite3
+import json
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from codegraph.errors import DatabaseError
+
+
+def get_database_path(project_root: str) -> str:
+    """Get the path to the SQLite database."""
+    return os.path.join(project_root, '.codegraph', 'codegraph.db')
+
+
+def remove_database_files(db_path: str) -> None:
+    """Remove the database file and its WAL/SHM sidecars."""
+    for suffix in ('', '-wal', '-shm'):
+        p = db_path + suffix
+        if os.path.exists(p):
+            os.remove(p)
+
+
+class DatabaseConnection:
+    """Manages a SQLite database connection."""
+
+    def __init__(self, db_path: str, db: sqlite3.Connection):
+        self._path = db_path
+        self._db = db
+        self._db.row_factory = sqlite3.Row
+
+        # ── Connection PRAGMAs (order matters) ──
+        # busy_timeout MUST be first — if another process holds a write lock,
+        # subsequent pragmas wait out the lock instead of throwing immediately.
+        self._db.execute("PRAGMA busy_timeout = 5000")
+        self._db.execute("PRAGMA foreign_keys = ON")
+        self._db.execute("PRAGMA journal_mode = WAL")      # Write-Ahead Logging
+        self._db.execute("PRAGMA synchronous = NORMAL")    # safe with WAL mode
+        self._db.execute("PRAGMA cache_size = -64000")     # 64 MB page cache
+        self._db.execute("PRAGMA temp_store = MEMORY")     # temp tables in memory
+        self._db.execute("PRAGMA mmap_size = 268435456")   # 256 MB memory-mapped I/O
+
+    @staticmethod
+    def initialize(db_path: str) -> 'DatabaseConnection':
+        """Create a new database with schema."""
+        from codegraph import directory as dir_mod
+        cg_dir = os.path.dirname(db_path)
+        os.makedirs(cg_dir, exist_ok=True)
+
+        # Create .gitignore inside .codegraph/ to prevent index from being tracked
+        gitignore_path = os.path.join(cg_dir, '.gitignore')
+        if not os.path.isfile(gitignore_path):
+            with open(gitignore_path, 'w') as f:
+                f.write('# CodeGraph index — auto-generated, not for version control\n*\n')
+
+        db = sqlite3.connect(db_path, check_same_thread=False)
+        conn = DatabaseConnection(db_path, db)
+
+        # Run schema
+        schema_path = os.path.join(os.path.dirname(__file__), 'schema.sql')
+        if os.path.isfile(schema_path):
+            with open(schema_path, 'r') as f:
+                schema = f.read()
+            db.executescript(schema)
+
+        # Set initial schema version
+        conn.set_metadata('schema_version', '1')
+
+        return conn
+
+    @staticmethod
+    def open(db_path: str) -> 'DatabaseConnection':
+        """Open an existing database."""
+        if not os.path.isfile(db_path):
+            raise DatabaseError(f'Database not found: {db_path}')
+        db = sqlite3.connect(db_path, check_same_thread=False)
+        return DatabaseConnection(db_path, db)
+
+    def get_db(self) -> sqlite3.Connection:
+        """Get the raw SQLite connection."""
+        return self._db
+
+    def get_path(self) -> str:
+        """Get the database file path."""
+        return self._path
+
+    def close(self) -> None:
+        """Close the database connection."""
+        try:
+            self._db.close()
+        except Exception:
+            pass
+
+    def run_maintenance(self) -> None:
+        """Run maintenance operations after bulk writes."""
+        try:
+            self._db.execute("PRAGMA optimize")
+        except Exception:
+            pass
+        try:
+            # Fold pending WAL pages back into the main DB file
+            self._db.execute("PRAGMA wal_checkpoint(PASSIVE)")
+        except Exception:
+            pass
+
+    # =========================================================================
+    # Metadata
+    # =========================================================================
+
+    def get_metadata(self, key: str) -> Optional[str]:
+        """Get a metadata value."""
+        try:
+            cur = self._db.execute(
+                'SELECT value FROM project_metadata WHERE key = ?', (key,)
+            )
+            row = cur.fetchone()
+            return row['value'] if row else None
+        except Exception:
+            return None
+
+    def set_metadata(self, key: str, value: str) -> None:
+        """Set a metadata value."""
+        try:
+            self._db.execute(
+                '''INSERT OR REPLACE INTO project_metadata (key, value, updated_at)
+                   VALUES (?, ?, ?)''',
+                (key, value, int(time.time() * 1000))
+            )
+        except Exception:
+            pass
+
+    def get_journal_mode(self) -> Optional[str]:
+        """Get the effective journal mode."""
+        try:
+            cur = self._db.execute("PRAGMA journal_mode")
+            row = cur.fetchone()
+            return row[0] if row else None
+        except Exception:
+            return None
+
+    def run_migrations(self) -> None:
+        """Run any pending database migrations."""
+        # Check current version
+        current = self.get_metadata('schema_version')
+        if current is None:
+            current = '1'
+
+        version = int(current)
+
+        # Apply migrations as needed
+        # Version 1 is the initial schema
+        # Future migrations will be added here
+
+        if version > 1:
+            self.set_metadata('schema_version', str(version))