"""Interactive D3.js graph visualization for code knowledge graphs.
Exports graph data to JSON and generates a self-contained HTML file with
a force-directed D3.js visualization. Dark theme, zoomable, draggable,
with collapsible file clusters, tooltips, legend, and stats bar.
Supports multiple rendering modes for large graphs:
- ``full`` — render every node (default, current behavior)
- ``community`` — aggregate by community; double-click to drill down
- ``file`` — aggregate by file; each file is a node
- ``auto`` — choose community mode when node count exceeds threshold
"""
from __future__ import annotations
import json
import logging
import sqlite3
from collections import Counter, defaultdict
from dataclasses import asdict
from pathlib import Path
from .graph import GraphStore, edge_to_dict, node_to_dict
logger = logging.getLogger(__name__)
def _build_name_index(
nodes: list[dict], seen_qn: set[str]
) -> dict[str, list[str]]:
"""Build a mapping from short/module-style names to qualified names.
Returns ``{short_name: [qualified_name, ...]}``.
"""
index: dict[str, list[str]] = {}
def _add(key: str, qn: str) -> None:
index.setdefault(key, []).append(qn)
for n in nodes:
qn = n["qualified_name"]
_add(n["name"], qn)
# Index by "file::name" suffix (e.g. "cli.py::main")
if "::" in qn:
_add(qn.rsplit("/", 1)[-1], qn)
# Index by module-style path (e.g. "merit.cli" or "merit.cli.main")
fp = n.get("file_path", "")
if fp:
mod = fp.replace("/", ".").replace(".py", "")
if n["kind"] == "File":
_add(mod, qn)
# Index by every path suffix so C/C++ bare includes resolve.
# e.g. "/abs/libs/trading/Foo.hpp" is also indexed as
# "Foo.hpp", "trading/Foo.hpp", "libs/trading/Foo.hpp", …
parts = fp.replace("\\", "/").split("/")
for i in range(len(parts)):
suffix = "/".join(parts[i:])
if suffix:
_add(suffix, qn)
else:
_add(mod + "." + n["name"], qn)
return index
def _resolve_target(
target: str,
source: str,
seen_qn: set[str],
name_index: dict[str, list[str]],
) -> str | None:
"""Try to resolve an unqualified edge target to a full qualified name.
Returns the resolved qualified name, or None if unresolvable.
"""
# Already fully qualified
if target in seen_qn:
return target
candidates = name_index.get(target)
if not candidates:
return None
if len(candidates) == 1:
return candidates[0]
# Disambiguate: prefer node in the same file as the source
src_file = source.split("::")[0] if "::" in source else source
same_file = [c for c in candidates if c.startswith(src_file)]
if len(same_file) == 1:
return same_file[0]
# Prefer node in the same top-level directory
src_parts = src_file.rsplit("/", 1)[0] if "/" in src_file else ""
same_dir = [c for c in candidates if c.startswith(src_parts)]
if len(same_dir) == 1:
return same_dir[0]
# Ambiguous — pick first match rather than dropping the edge
return candidates[0]
def export_graph_data(store: GraphStore) -> dict:
"""Export all graph nodes and edges as a JSON-serializable dict.
Returns ``{"nodes": [...], "edges": [...], "stats": {...},
"flows": [...], "communities": [...]}``.
"""
nodes = []
seen_qn: set[str] = set()
# Preload community_id mapping from DB (column may not exist in old schemas)
community_map = store.get_all_community_ids()
for file_path in store.get_all_files():
for gnode in store.get_nodes_by_file(file_path):
if gnode.qualified_name in seen_qn:
continue
seen_qn.add(gnode.qualified_name)
d = node_to_dict(gnode)
d["params"] = gnode.params
d["return_type"] = gnode.return_type
d["community_id"] = community_map.get(gnode.qualified_name)
nodes.append(d)
name_index = _build_name_index(nodes, seen_qn)
all_edges = [edge_to_dict(e) for e in store.get_all_edges()]
# Resolve short/unqualified edge targets to full qualified names,
# then drop edges that still can't be resolved (external/stdlib calls).
edges = []
for e in all_edges:
src = _resolve_target(e["source"], e["source"], seen_qn, name_index)
tgt = _resolve_target(e["target"], e["source"], seen_qn, name_index)
if src and tgt:
e["source"] = src
e["target"] = tgt
edges.append(e)
stats = store.get_stats()
# Include flows (graceful fallback if table doesn't exist)
try:
from code_review_graph.flows import get_flows
flows = get_flows(store, limit=100)
except (ImportError, sqlite3.OperationalError) as exc:
logger.debug("flows unavailable for export: %s", exc)
flows = []
# Include communities (graceful fallback if table doesn't exist)
try:
from code_review_graph.communities import get_communities
communities = get_communities(store)
except (ImportError, sqlite3.OperationalError) as exc:
logger.debug("communities unavailable for export: %s", exc)
communities = []
return {
"nodes": nodes,
"edges": edges,
"stats": asdict(stats),
"flows": flows,
"communities": communities,
}
def _aggregate_community(data: dict) -> dict:
"""Aggregate full graph data into community-level super-nodes.
Each community becomes a single node sized by member count.
Edges between super-nodes represent the count of cross-community edges.
Returns a new dict with the same schema as *data* but fewer nodes/edges.
Also returns per-community detail data for drill-down rendering.
"""
communities = data.get("communities") or []
nodes = data["nodes"]
edges = data["edges"]
# Build mapping: qualified_name -> community_id
qn_to_cid: dict[str, int] = {}
for c in communities:
for qn in c.get("members", []):
qn_to_cid[qn] = c["id"]
# Also use node-level community_id for nodes not in community member lists
for n in nodes:
if n.get("community_id") is not None and n["qualified_name"] not in qn_to_cid:
qn_to_cid[n["qualified_name"]] = n["community_id"]
# Assign uncategorized nodes to a synthetic community id = -1
uncategorized_members: list[str] = []
for n in nodes:
if n["qualified_name"] not in qn_to_cid:
qn_to_cid[n["qualified_name"]] = -1
uncategorized_members.append(n["qualified_name"])
# Build community info map (including the synthetic uncategorized one)
cid_info: dict[int, dict] = {}
for c in communities:
cid_info[c["id"]] = c
if uncategorized_members:
cid_info[-1] = {
"id": -1,
"name": "Uncategorized",
"size": len(uncategorized_members),
"members": uncategorized_members,
"dominant_language": "",
"description": "Nodes not assigned to any community",
"cohesion": 0,
"level": 0,
}
# Build super-nodes (one per community)
super_nodes = []
for cid, info in cid_info.items():
size = info.get("size", len(info.get("members", [])))
if size == 0:
continue
super_nodes.append({
"qualified_name": f"__community__{cid}",
"name": info.get("name", f"Community {cid}"),
"kind": "Community",
"file_path": "",
"line_start": None,
"line_end": None,
"language": info.get("dominant_language", ""),
"community_id": cid,
"member_count": size,
"description": info.get("description", ""),
"id": cid,
})
# Build super-edges: aggregate cross-community edges
cross_edge_counts: Counter[tuple[int, int]] = Counter()
for e in edges:
src_cid = qn_to_cid.get(e["source"])
tgt_cid = qn_to_cid.get(e["target"])
if src_cid is not None and tgt_cid is not None and src_cid != tgt_cid:
pair = (min(src_cid, tgt_cid), max(src_cid, tgt_cid))
cross_edge_counts[pair] += 1
super_edges = []
for (c1, c2), count in cross_edge_counts.items():
super_edges.append({
"source": f"__community__{c1}",
"target": f"__community__{c2}",
"kind": "CROSS_COMMUNITY",
"weight": count,
})
# Build per-community detail data for drill-down
community_details: dict[int, dict] = {}
cid_members_set: dict[int, set[str]] = defaultdict(set)
for qn, cid in qn_to_cid.items():
cid_members_set[cid].add(qn)
for cid, member_qns in cid_members_set.items():
detail_nodes = [n for n in nodes if n["qualified_name"] in member_qns]
detail_edges = [
e for e in edges
if e["source"] in member_qns and e["target"] in member_qns
]
community_details[cid] = {
"nodes": detail_nodes,
"edges": detail_edges,
}
return {
"nodes": super_nodes,
"edges": super_edges,
"stats": data["stats"],
"flows": data.get("flows", []),
"communities": communities,
"mode": "community",
"community_details": {
str(k): v for k, v in community_details.items()
},
}
def _aggregate_file(data: dict) -> dict:
"""Aggregate full graph data into file-level nodes.
Each file becomes a node sized by symbol count.
Edges between files represent aggregated cross-file dependencies.
"""
nodes = data["nodes"]
edges = data["edges"]
# Count symbols per file
file_symbol_count: Counter[str] = Counter()
qn_to_file: dict[str, str] = {}
file_languages: dict[str, str] = {}
for n in nodes:
fp = n.get("file_path", "")
if not fp:
continue
qn_to_file[n["qualified_name"]] = fp
if n["kind"] != "File":
file_symbol_count[fp] += 1
else:
file_symbol_count.setdefault(fp, 0)
if n.get("language"):
file_languages[fp] = n["language"]
# Build file nodes
file_nodes = []
for fp, count in file_symbol_count.items():
parts = fp.replace("\\", "/").split("/")
short = parts[-1] if parts else fp
parent = parts[-2] if len(parts) >= 2 else ""
label = f"{parent}/{short}" if parent else short
# Recover community_id from the majority of symbols in this file
cid = None
for n in nodes:
if n.get("file_path") == fp and n.get("community_id") is not None:
cid = n["community_id"]
break
file_nodes.append({
"qualified_name": fp,
"name": label,
"kind": "File",
"file_path": fp,
"line_start": None,
"line_end": None,
"language": file_languages.get(fp, ""),
"community_id": cid,
"symbol_count": count,
})
# Aggregate cross-file edges
cross_file_counts: Counter[tuple[str, str]] = Counter()
for e in edges:
src_fp = qn_to_file.get(e["source"])
tgt_fp = qn_to_file.get(e["target"])
if src_fp and tgt_fp and src_fp != tgt_fp:
pair = (src_fp, tgt_fp)
cross_file_counts[pair] += 1
file_edges = []
for (f1, f2), count in cross_file_counts.items():
file_edges.append({
"source": f1,
"target": f2,
"kind": "DEPENDS_ON",
"weight": count,
})
return {
"nodes": file_nodes,
"edges": file_edges,
"stats": data["stats"],
"flows": data.get("flows", []),
"communities": data.get("communities", []),
"mode": "file",
}
def generate_html(
store: GraphStore,
output_path: str | Path,
mode: str = "auto",
max_full_nodes: int = 3000,
) -> Path:
"""Generate a self-contained interactive HTML visualization.
Args:
store: The GraphStore to read graph data from.
output_path: Path for the output HTML file.
mode: Rendering mode — ``"auto"``, ``"full"``, ``"community"``,
or ``"file"``. ``"auto"`` switches to ``"community"`` when
the node count exceeds *max_full_nodes*.
max_full_nodes: Threshold for auto-switching to community mode.
Writes the HTML file to *output_path* and returns the resolved Path.
"""
output_path = Path(output_path)
stats = store.get_stats()
if stats.total_nodes > 50000:
logger.warning(
"Graph has %d nodes — visualization may be slow. "
"Consider filtering by file pattern.", stats.total_nodes,
)
data = export_graph_data(store)
# Determine effective mode
effective_mode = mode
if effective_mode == "auto":
effective_mode = (
"community" if stats.total_nodes > max_full_nodes else "full"
)
if effective_mode == "community":
# Keep full data available for drill-down; aggregate for top-level
agg = _aggregate_community(data)
# Escape inside JSON to prevent premature tag closure
data_json = json.dumps(agg, default=str).replace("", "<\\/")
html = _AGGREGATED_HTML_TEMPLATE.replace("__GRAPH_DATA__", data_json)
elif effective_mode == "file":
agg = _aggregate_file(data)
data_json = json.dumps(agg, default=str).replace("", "<\\/")
html = _AGGREGATED_HTML_TEMPLATE.replace("__GRAPH_DATA__", data_json)
else:
# full mode — original behavior
data_json = json.dumps(data, default=str).replace("", "<\\/")
html = _HTML_TEMPLATE.replace("__GRAPH_DATA__", data_json)
output_path.write_text(html, encoding="utf-8")
return output_path
# ---------------------------------------------------------------------------
# Full D3.js interactive HTML template
# ---------------------------------------------------------------------------
# Template lives in this file for zero-dependency packaging (no external files
# to locate at runtime). The E501 suppression for this module is configured via
# pyproject.toml per-file-ignores for this reason.
_HTML_TEMPLATE = r"""
Click edge types in the legend to toggle visibility
Keyboard Shortcuts
/
Focus search
?
Toggle this help
Esc
Close panel / search / help
Enter / Space
Activate focused node
Arrow keys
Navigate between nodes
Click anywhere outside to dismiss
Laying out graph…
🔍
No nodes to display
The graph is empty. Run code-review-graph build to index your codebase, then regenerate the visualization.
"""
# ---------------------------------------------------------------------------
# Aggregated-mode HTML template (community / file)
# ---------------------------------------------------------------------------
# Supports community super-nodes with drill-down (double-click) and a Back
# button to return to the overview.
# NOTE: innerHTML / insertAdjacentHTML usage below mirrors the original
# _HTML_TEMPLATE and is safe because all interpolated values pass through
# escH() which escapes &, <, >, ", ', and backtick characters.
_AGGREGATED_HTML_TEMPLATE = r"""
Code Review Graph (Aggregated)