diff --git a/ast_rag/cli.py b/ast_rag/cli.py index e22741b..0d32f51 100644 --- a/ast_rag/cli.py +++ b/ast_rag/cli.py @@ -1806,5 +1806,116 @@ def cache_stats( console.rule("[bold green]Done[/bold green]") +# --------------------------------------------------------------------------- + + +def _collect_index_stats(driver) -> dict: + """Gather node/edge/language/file counts from the graph. + + Edges are stored as generic ``:EDGE`` relationships carrying the semantic + type in a ``kind`` property (see ``repositories.queries.batch_upsert_edges``), + so edge types are grouped by that property rather than by relationship type. + """ + stats: dict = { + "nodes": {"total": 0, "by_kind": {}}, + "edges": {"total": 0, "by_kind": {}}, + "languages": {}, + "files": 0, + } + + with driver.session() as session: + # CurrentVersion is the graph's bookkeeping pointer, not a code node. + for record in session.run( + "MATCH (n) WHERE n.id IS NOT NULL AND NOT n:CurrentVersion " + "RETURN labels(n)[0] AS kind, count(*) AS n ORDER BY n DESC" + ): + stats["nodes"]["by_kind"][record["kind"] or "unlabelled"] = record["n"] + stats["nodes"]["total"] = sum(stats["nodes"]["by_kind"].values()) + + for record in session.run( + "MATCH ()-[r]->() RETURN coalesce(r.kind, type(r)) AS kind, " + "count(*) AS n ORDER BY n DESC" + ): + stats["edges"]["by_kind"][record["kind"]] = record["n"] + stats["edges"]["total"] = sum(stats["edges"]["by_kind"].values()) + + for record in session.run( + "MATCH (n) WHERE n.lang IS NOT NULL " + "RETURN n.lang AS lang, count(*) AS n ORDER BY n DESC" + ): + stats["languages"][record["lang"]] = record["n"] + + record = session.run( + "MATCH (n) WHERE n.file_path IS NOT NULL RETURN count(DISTINCT n.file_path) AS files" + ).single() + stats["files"] = record["files"] if record else 0 + + return stats + + +@app.command("stats") +def stats( + config: Optional[str] = typer.Option(None, "--config", "-c", help="Path to config JSON"), + as_json: bool = typer.Option(False, "--json", help="Emit JSON instead of a table"), + verbose: bool = typer.Option(False, "--verbose", "-v"), +) -> None: + """ + Show statistics about the indexed codebase. + + Reports node counts by kind, edge counts by type, the language + distribution and the number of indexed files. + + Examples: + + ast-rag stats + ast-rag stats --json + ast-rag stats --config ast_rag_config.json + """ + if verbose: + logging.basicConfig(level=logging.DEBUG) + else: + logging.basicConfig(level=logging.WARNING) + + cfg = _load_config(config) + driver = create_driver(cfg.neo4j) + try: + data = _collect_index_stats(driver) + finally: + driver.close() + + if as_json: + print(json.dumps(data, indent=2)) + return + + if data["nodes"]["total"] == 0: + console.print("[yellow]No indexed nodes found. Run `ast-rag init ` first.[/yellow]") + return + + console.rule("[bold blue]AST-RAG Index Statistics[/bold blue]") + console.print(f" Files indexed: {data['files']:,}") + + console.print() + console.print(f"[bold]Nodes[/bold] ({data['nodes']['total']:,} total)") + for kind, count in data["nodes"]["by_kind"].items(): + console.print(f" {kind:<20} {count:>8,}") + + console.print() + console.print(f"[bold]Edges[/bold] ({data['edges']['total']:,} total)") + if data["edges"]["by_kind"]: + for kind, count in data["edges"]["by_kind"].items(): + console.print(f" {kind:<20} {count:>8,}") + else: + console.print(" [dim]none[/dim]") + + console.print() + console.print("[bold]Languages[/bold]") + if data["languages"]: + total = sum(data["languages"].values()) or 1 + for lang, count in data["languages"].items(): + console.print(f" {lang:<20} {count:>8,} ({count / total:.0%})") + else: + console.print(" [dim]none[/dim]") + + if __name__ == "__main__": app() diff --git a/tests/test_stats_command.py b/tests/test_stats_command.py new file mode 100644 index 0000000..10899e7 --- /dev/null +++ b/tests/test_stats_command.py @@ -0,0 +1,98 @@ +"""Tests for the `ast-rag stats` command (issue #13).""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +from typer.testing import CliRunner + +from ast_rag.cli import _collect_index_stats, app + +runner = CliRunner() + + +class _Rec(dict): + """Minimal stand-in for a neo4j Record.""" + + def __getitem__(self, key): + return super().__getitem__(key) + + +def _driver_returning(node_rows, edge_rows, lang_rows, files): + """Build a mock driver whose session.run dispatches on the query text.""" + session = MagicMock() + + def run(query, **_kwargs): + result = MagicMock() + if "labels(n)[0]" in query: + return iter([_Rec(kind=k, n=n) for k, n in node_rows]) + if "coalesce(r.kind" in query: + return iter([_Rec(kind=k, n=n) for k, n in edge_rows]) + if "n.lang" in query: + return iter([_Rec(lang=lang, n=n) for lang, n in lang_rows]) + result.single.return_value = _Rec(files=files) + return result + + session.run.side_effect = run + driver = MagicMock() + driver.session.return_value.__enter__.return_value = session + return driver + + +def test_collect_index_stats_shapes_and_totals(): + driver = _driver_returning( + node_rows=[("Method", 368), ("Class", 100)], + edge_rows=[("CONTAINS_METHOD", 392), ("CALLS", 150)], + lang_rows=[("java", 468)], + files=103, + ) + stats = _collect_index_stats(driver) + + assert stats["files"] == 103 + assert stats["nodes"]["total"] == 468 + assert stats["nodes"]["by_kind"] == {"Method": 368, "Class": 100} + assert stats["edges"]["total"] == 542 + assert stats["edges"]["by_kind"]["CALLS"] == 150 + assert stats["languages"] == {"java": 468} + + +def test_edges_grouped_by_kind_property_not_relationship_type(): + """Edges are stored as :EDGE with the semantic type in `kind`.""" + driver = _driver_returning([], [], [], 0) + _collect_index_stats(driver) + queries = [ + c.args[0] for c in driver.session.return_value.__enter__.return_value.run.call_args_list + ] + edge_query = next(q for q in queries if "]->()" in q and "count(*)" in q) + assert "r.kind" in edge_query, "edge stats must group by the kind property" + + +def test_bookkeeping_node_excluded(): + driver = _driver_returning([], [], [], 0) + _collect_index_stats(driver) + queries = [ + c.args[0] for c in driver.session.return_value.__enter__.return_value.run.call_args_list + ] + node_query = next(q for q in queries if "labels(n)[0]" in q) + assert "CurrentVersion" in node_query + + +def test_empty_index_reports_hint(monkeypatch): + monkeypatch.setattr("ast_rag.cli.create_driver", lambda _cfg: _driver_returning([], [], [], 0)) + result = runner.invoke(app, ["stats"]) + assert result.exit_code == 0 + assert "No indexed nodes found" in result.stdout + + +def test_json_output_is_valid_json(monkeypatch): + monkeypatch.setattr( + "ast_rag.cli.create_driver", + lambda _cfg: _driver_returning([("Method", 5)], [("CALLS", 2)], [("python", 5)], 3), + ) + result = runner.invoke(app, ["stats", "--json"]) + assert result.exit_code == 0 + payload = json.loads(result.stdout) + assert payload["nodes"]["by_kind"] == {"Method": 5} + assert payload["edges"]["by_kind"] == {"CALLS": 2} + assert payload["files"] == 3