"""Module overview diagram generation using Mermaid."""
from __future__ import annotations
from collections import defaultdict, Counter
from pathlib import Path
from local_deepwiki.models import IndexStatus
from ._utils import sanitize_mermaid_name
def generate_module_overview(
index_status: IndexStatus,
show_file_counts: bool = True,
) -> str | None:
"""Generate a high-level module overview diagram.
Shows package structure with subgraphs for major directories.
Args:
index_status: Index status with file information.
show_file_counts: Whether to show file counts in nodes.
Returns:
Mermaid diagram string, or None if not enough structure.
"""
if not index_status.files:
return None
# Known artifact directories to exclude even if they slipped into the index
_ARTIFACT_DIRS = frozenset(
{
"htmlcov",
"coverage",
".pytest_cache",
".mypy_cache",
".ruff_cache",
".tox",
".nox",
".eggs",
}
)
# Group files by top-level directory
directories: defaultdict[str, Counter[str]] = defaultdict(Counter)
for file_info in index_status.files:
parts = list(Path(file_info.path).parts)
if len(parts) < 2:
continue
# Skip artifact directories
if any(p in _ARTIFACT_DIRS for p in parts):
continue
top_dir = parts[0]
if top_dir in ("src", "lib", "pkg"):
if len(parts) > 1:
top_dir = parts[1]
parts = parts[1:]
if len(parts) > 1:
directories[top_dir][parts[1]] += 1
else:
directories[top_dir]["_root"] += 1
if not directories:
return None
# Build diagram
lines = ["```mermaid", "graph TB"]
for top_dir, subdirs in sorted(directories.items()):
safe_dir = sanitize_mermaid_name(top_dir)
total_files = sum(subdirs.values())
if len(subdirs) > 1 and "_root" not in subdirs:
# Create subgraph for directories with multiple subdirs
lines.append(f" subgraph {safe_dir}[{top_dir}]")
for subdir, count in sorted(subdirs.items()):
if subdir != "_root":
safe_sub = sanitize_mermaid_name(f"{top_dir}_{subdir}")
label = f"{subdir}"
if show_file_counts:
label += f" ({count})"
lines.append(f" {safe_sub}[{label}]")
lines.append(" end")
else:
# Single node for simple directories
label = top_dir
if show_file_counts:
label += f" ({total_files})"
lines.append(f" {safe_dir}[{label}]")
lines.append("```")
return "\n".join(lines)
