niishantth/codegraph

1"""CLI interface for CodeGraph."""2 3from __future__ import annotations4 5import json6import sys7from pathlib import Path8 9import click10from rich.console import Console11from rich.table import Table12from rich.tree import Tree13 14from codegraph.parser import CodebaseScanner15from codegraph.graph_builder import GraphBuilder16from codegraph.query_engine import QueryEngine17from codegraph.models import EntityType18 19console = Console()20 21 22def _entity_icon(e_type: EntityType) -> str:23    return {24        EntityType.FILE: "📄",25        EntityType.MODULE: "📦",26        EntityType.CLASS: "🏛",27        EntityType.FUNCTION: "🔧",28        EntityType.METHOD: "🔩",29        EntityType.VARIABLE: "📌",30        EntityType.IMPORT: "⬇",31        EntityType.CALL: "📞",32        EntityType.ERROR: "❌",33    }.get(e_type, "❓")34 35 36@click.group()37@click.version_option(version="0.1.0")38def cli() -> None:39    """CodeGraph – Knowledge graph of your entire codebase."""40    pass41 42 43@cli.command()44@click.argument("path", type=click.Path(exists=True, file_okay=False))45@click.option("--output", "-o", type=click.Path(), help="JSON file to save the graph.")46@click.option("--ignore", "-i", multiple=True, help="Directory names to ignore.")47def build(path: str, output: str | None, ignore: tuple[str, ...]) -> None:48    """Scan a codebase and build the knowledge graph."""49    scanner = CodebaseScanner(path, ignore_dirs=set(ignore) if ignore else None)50    with console.status("[bold green]Scanning codebase..."):51        entities = scanner.scan()52    console.print(f"[green]Parsed {len(entities)} entities from {len(scanner.files)} files.[/green]")53 54    builder = GraphBuilder()55    with console.status("[bold green]Building graph..."):56        graph = builder.build_from_entities(entities)57    console.print(f"[green]Graph has {graph.number_of_nodes()} nodes and {graph.number_of_edges()} edges.[/green]")58 59    if output:60        data = {61            "nodes": [builder.entity_map[n].model_dump() for n in graph.nodes()],62            "edges": [63                {"source": u, "target": v, **graph.edges[u, v]}64                for u, v in graph.edges()65            ],66        }67        Path(output).write_text(json.dumps(data, indent=2, default=str))68        console.print(f"[blue]Graph saved to {output}[/blue]")69 70    # Store builder in a temp file for subsequent commands71    _save_state(builder)72 73 74@cli.command()75@click.argument("name")76def find(name: str) -> None:77    """Find entities by name."""78    builder = _load_state()79    if builder is None:80        console.print("[red]No graph found. Run `codegraph build <path>` first.[/red]")81        sys.exit(1)82    engine = QueryEngine(builder)83    results = engine.find_entity(name)84    if not results:85        console.print(f"[yellow]No entities named '{name}' found.[/yellow]")86        return87    table = Table(title=f"Entities matching '{name}'")88    table.add_column("Type", style="cyan")89    table.add_column("Name", style="magenta")90    table.add_column("File", style="green")91    table.add_column("Line", style="yellow")92    for e in results:93        line = str(e.span.start.line) if e.span else "-"94        table.add_row(_entity_icon(e.type) + " " + e.type.value, e.name, e.file_path, line)95    console.print(table)96 97 98@cli.command()99@click.argument("file_path", type=click.Path(exists=True))100def outline(file_path: str) -> None:101    """Show the outline (top-level entities) of a file."""102    builder = _load_state()103    if builder is None:104        console.print("[red]No graph found. Run `codegraph build <path>` first.[/red]")105        sys.exit(1)106    engine = QueryEngine(builder)107    ents = engine.file_outline(str(Path(file_path).resolve()))108    if not ents:109        console.print(f"[yellow]No entities found in {file_path}.[/yellow]")110        return111    tree = Tree(f"📄 {file_path}")112    for e in ents:113        tree.add(f"{_entity_icon(e.type)} {e.type.value} [bold]{e.name}[/bold] (L{e.span.start.line if e.span else '?'})")114    console.print(tree)115 116 117@cli.command(name="map-error")118@click.argument("error_message")119@click.option("--file", "file_path", help="File path associated with the error.")120@click.option("--line", type=int, help="Line number associated with the error.")121@click.option("--column", type=int, help="Column number associated with the error.")122def map_error(error_message: str, file_path: str | None, line: int | None, column: int | None) -> None:123    """Map an error message to the knowledge graph."""124    builder = _load_state()125    if builder is None:126        console.print("[red]No graph found. Run `codegraph build <path>` first.[/red]")127        sys.exit(1)128    engine = QueryEngine(builder)129    mapping = engine.map_error(error_message, file_path=file_path, line=line, column=column)130 131    console.print(f"[bold]Error Type:[/bold] {mapping.error_type}")132    console.print(f"[bold]Message:[/bold] {mapping.error_message}")133    if mapping.matched_entity:134        e = mapping.matched_entity135        console.print(f"[bold]Mapped to:[/bold] {_entity_icon(e.type)} {e.type.value} [cyan]{e.name}[/cyan] in {e.file_path}:{e.span.start.line if e.span else '?'}")136    else:137        console.print("[yellow]Could not map error to a specific entity.[/yellow]")138 139    if mapping.related_entities:140        console.print("[bold]Related entities:[/bold]")141        for r in mapping.related_entities:142            console.print(f"  - {_entity_icon(r.type)} {r.type.value} {r.name} ({r.file_path}:{r.span.start.line if r.span else '?'})")143 144    if mapping.suggested_fix:145        console.print(f"[bold green]Suggested fix:[/bold green] {mapping.suggested_fix}")146 147 148@cli.command()149@click.argument("query")150def search(query: str) -> None:151    """Fuzzy search across the graph."""152    builder = _load_state()153    if builder is None:154        console.print("[red]No graph found. Run `codegraph build <path>` first.[/red]")155        sys.exit(1)156    engine = QueryEngine(builder)157    results = engine.search(query)158    if not results:159        console.print(f"[yellow]No results for '{query}'.[/yellow]")160        return161    table = Table(title=f"Search results for '{query}'")162    table.add_column("Type", style="cyan")163    table.add_column("Name", style="magenta")164    table.add_column("File", style="green")165    table.add_column("Line", style="yellow")166    for e in results[:50]:167        line = str(e.span.start.line) if e.span else "-"168        table.add_row(_entity_icon(e.type) + " " + e.type.value, e.name, e.file_path, line)169    console.print(table)170 171 172# ------------------------------------------------------------------173# State persistence helpers174# ------------------------------------------------------------------175 176_STATE_PATH = Path("/tmp/codegraph_state.json")177 178 179def _save_state(builder: GraphBuilder) -> None:180    data = {181        "nodes": [builder.entity_map[n].model_dump() for n in builder.graph.nodes()],182        "edges": [183            {"source": u, "target": v, **builder.graph.edges[u, v]}184            for u, v in builder.graph.edges()185        ],186    }187    _STATE_PATH.write_text(json.dumps(data, default=str))188 189 190def _load_state() -> GraphBuilder | None:191    if not _STATE_PATH.exists():192        return None193    raw = json.loads(_STATE_PATH.read_text())194    from codegraph.models import Entity, Relationship, RelationType195 196    builder = GraphBuilder()197    entities = [Entity(**n) for n in raw["nodes"]]198    builder.add_entities(entities)199    rels = [200        Relationship(source=e["source"], target=e["target"], type=RelationType(e.get("type", "uses")), meta={k: v for k, v in e.items() if k not in ("source", "target", "type")})201        for e in raw["edges"]202    ]203    builder.add_relationships(rels)204    return builder205 206 207if __name__ == "__main__":208    cli()209