generate.py

#!/usr/bin/env python3
"""Generate data.json — an exact filesystem tree index for any repository.

Usage:
  # Point at any repo (writes workspaces/<slug>/data.json; other repos stay put)
  python generate.py --root /path/to/repo

  # Embedded as <repo>/codebase-graph/ — scans the parent repo
  python generate.py

  # Tool repo with an attached clone in ./repo
  python generate.py
  # (auto-detects ./repo if present)

  # CI / freshness check
  python generate.py --root /path/to/repo --check
"""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, timezone
from pathlib import Path

# Names skipped at any depth (tooling / caches / VCS).
EXCLUDE_NAMES = {
    ".git",
    ".hg",
    ".svn",
    ".venv",
    "venv",
    ".ruff_cache",
    ".mypy_cache",
    ".pytest_cache",
    ".tox",
    ".mimicflow",
    ".claude",
    ".cursor",
    ".idea",
    ".vs",
    ".DS_Store",
    "__pycache__",
    "node_modules",
    "target",
    "dist",
    "build",
    "coverage",
    # Avoid indexing this tool when it lives inside the scanned tree.
    "codebase-graph",
    "workspaces",
}

# Optional path prefixes skipped (posix, relative to scan root). Override with --exclude.
DEFAULT_EXCLUDE_PREFIXES: tuple[str, ...] = (
    "tests/wpt",  # common huge upstream checkout; harmless elsewhere
)


def looks_like_repo(path: Path) -> bool:
    return (path / ".git").exists() or (path / ".git").is_file()


def detect_root(tool_dir: Path) -> Path:
    """Resolve which directory to scan when --root is omitted."""
    attached = tool_dir / "repo"
    if attached.is_dir() and looks_like_repo(attached):
        return attached

    parent = tool_dir.parent
    # Copied/cloned into an existing project as <project>/codebase-graph/
    if tool_dir.name == "codebase-graph" and looks_like_repo(parent):
        return parent

    # Standalone tool checkout: prefer cwd if it is a repo, else tool dir itself
    cwd = Path.cwd().resolve()
    if looks_like_repo(cwd) and cwd != tool_dir:
        return cwd

    return tool_dir


def should_exclude(rel_posix: str, name: str, prefixes: tuple[str, ...]) -> bool:
    if name in EXCLUDE_NAMES:
        return True
    for prefix in prefixes:
        if not prefix:
            continue
        if rel_posix == prefix or rel_posix.startswith(prefix.rstrip("/") + "/"):
            return True
    return False


def count_file_loc(path: Path) -> int:
    """Newline count for text-ish files; 0 for binary / unreadable."""
    try:
        data = path.read_bytes()
    except OSError:
        return 0
    if not data:
        return 0
    sample = data[:8192]
    if b"\0" in sample:
        return 0
    # Reject high control-char ratio as binary
    ctrl = sum(1 for b in sample if b < 9 or (13 < b < 32))
    if ctrl > max(8, len(sample) // 20):
        return 0
    return data.count(b"\n") + (0 if data.endswith(b"\n") else 1)


def scan_dir(
    abs_path: Path,
    rel_posix: str,
    stats: dict[str, int],
    prefixes: tuple[str, ...],
    root_name: str,
) -> dict:
    children: list[dict] = []
    try:
        entries = list(abs_path.iterdir())
    except OSError as exc:
        print(f"warning: cannot read {rel_posix or '.'}: {exc}", file=sys.stderr)
        entries = []

    entries.sort(key=lambda p: (not p.is_dir(), p.name.lower()))
    loc_total = 0

    for entry in entries:
        name = entry.name
        child_rel = f"{rel_posix}/{name}" if rel_posix else name
        if should_exclude(child_rel, name, prefixes):
            continue
        if entry.is_dir() and not entry.is_symlink():
            stats["dirs"] += 1
            child = scan_dir(entry, child_rel, stats, prefixes, root_name)
            loc_total += int(child.get("loc") or 0)
            children.append(child)
        elif entry.is_file() or entry.is_symlink():
            if entry.is_symlink() and entry.is_dir():
                continue
            stats["files"] += 1
            loc = count_file_loc(entry) if entry.is_file() else 0
            loc_total += loc
            stats["lines"] = stats.get("lines", 0) + loc
            children.append({"name": name, "kind": "file", "loc": loc})

    node_name = abs_path.name if rel_posix else root_name
    return {"name": node_name, "kind": "dir", "loc": loc_total, "children": children}


def build_payload(root: Path, prefixes: tuple[str, ...]) -> dict:
    root_name = root.name or "repo"
    stats = {"dirs": 1, "files": 0, "lines": 0}
    tree = scan_dir(root, "", stats, prefixes, root_name)
    return {
        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "root_name": root_name,
        "root_path": str(root),
        "stats": stats,
        "tree": tree,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--root",
        type=Path,
        default=None,
        help="Directory to index (default: auto-detect ./repo, parent project, or cwd)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Output JSON path (default: workspaces/<slug>/data.json)",
    )
    parser.add_argument(
        "--exclude",
        action="append",
        default=[],
        metavar="PREFIX",
        help="Extra repo-relative path prefix to skip (repeatable). "
        "Use --exclude '' with --no-default-excludes to clear defaults.",
    )
    parser.add_argument(
        "--no-default-excludes",
        action="store_true",
        help="Do not apply DEFAULT_EXCLUDE_PREFIXES (still skips EXCLUDE_NAMES)",
    )
    parser.add_argument(
        "--check",
        action="store_true",
        help="Exit 1 if existing data.json tree/stats differ (ignores generated_at)",
    )
    args = parser.parse_args()

    import workspace as workspace_mod

    workspace_mod.migrate_legacy()
    tool_dir = Path(__file__).resolve().parent
    root = (args.root or detect_root(tool_dir)).resolve()
    if args.output:
        output = args.output.resolve()
    else:
        output = workspace_mod.paths_for(root).data_path

    if not root.is_dir():
        print(f"error: root is not a directory: {root}", file=sys.stderr)
        return 1

    prefixes: tuple[str, ...]
    if args.no_default_excludes:
        prefixes = tuple(p for p in args.exclude if p is not None and p != "")
    else:
        prefixes = DEFAULT_EXCLUDE_PREFIXES + tuple(args.exclude)

    payload = build_payload(root, prefixes)

    if args.check:
        if not output.is_file():
            print(f"error: missing {output} — run without --check to generate it", file=sys.stderr)
            return 1
        existing = json.loads(output.read_text(encoding="utf-8"))
        left = {
            "root_name": existing.get("root_name"),
            "stats": existing.get("stats"),
            "tree": existing.get("tree"),
        }
        right = {
            "root_name": payload["root_name"],
            "stats": payload["stats"],
            "tree": payload["tree"],
        }
        if left != right:
            print("error: data.json is stale; run:", file=sys.stderr)
            print(f"  python {Path(__file__).name} --root {root}", file=sys.stderr)
            return 1
        print("data.json is up to date")
        return 0

    output.parent.mkdir(parents=True, exist_ok=True)
    text = json.dumps(payload, indent=2, ensure_ascii=False) + "\n"
    output.write_text(text, encoding="utf-8")
    if not args.output:
        workspace_mod.activate(root)
    try:
        shown = output.relative_to(Path.cwd())
    except ValueError:
        shown = output
    print(f"wrote {shown} from {root} ({payload['stats']['dirs']} dirs, {payload['stats']['files']} files)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())