giant/scripts/bump_dataset_version.py

#!/usr/bin/env python3
"""Cut a new raw generation or processed schema version for the geant_steps
dataset tree (see scripts/migrate_geant_steps.py for the layout):

    raw/<kind>/<gen>/<detector>/shard-NNN.root
    processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet

`gen` bumps when the underlying ROOT changes (geometry/physics-list/macro).
`schema` bumps when the parquet export (steps_to_parquet.py or similar)
changes, and is scoped to its gen — a new gen always starts back at schema1.

Creates the new (empty) target directory and appends a dated, reasoned entry
to VERSIONS.md. Defaults to a dry run; pass --execute to apply.

Usage:
    bump_dataset_version.py bump-gen --kind steps --reason "switched EM physics list"
    bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column"
    bump_dataset_version.py status
"""

import argparse
import datetime as dt
import re
import subprocess
from pathlib import Path

GEN_RE = re.compile(r"^gen(\d+)$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")


def _max_index(parent: Path, pattern: re.Pattern) -> int:
    """Highest N across child dir names matching *pattern* (0 if none/missing)."""
    if not parent.is_dir():
        return 0
    best = 0
    for child in parent.iterdir():
        m = pattern.match(child.name)
        if m and child.is_dir():
            best = max(best, int(m.group(1)))
    return best


def _git_user_name() -> str | None:
    try:
        out = subprocess.run(
            ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
        )
    except OSError:
        return None
    name = out.stdout.strip()
    return name or None


def plan_bump_gen(
    root: Path, kind: str, reason: str, by: str | None, date: str
) -> tuple[list[Path], str]:
    """New gen tag is one past the highest seen under raw/ or processed/ for *kind*
    — checking both, since a gen can exist in one tree before the other catches up."""
    next_gen = (
        max(
            _max_index(root / "raw" / kind, GEN_RE),
            _max_index(root / "processed" / kind, GEN_RE),
        )
        + 1
    )
    gen_tag = f"gen{next_gen}"
    new_dirs = [
        root / "raw" / kind / gen_tag,
        root / "processed" / kind / gen_tag / "schema1",
    ]
    by_suffix = f" ({by})" if by else ""
    log_line = f"- `{gen_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
    return new_dirs, log_line


def plan_bump_schema(
    root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str
) -> tuple[list[Path], str]:
    if not GEN_RE.match(gen_tag):
        raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}")
    raw_gen_dir = root / "raw" / kind / gen_tag
    processed_gen_dir = root / "processed" / kind / gen_tag
    if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
        raise SystemExit(
            f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
        )
    next_schema = _max_index(processed_gen_dir, SCHEMA_RE) + 1
    schema_tag = f"schema{next_schema}"
    new_dirs = [processed_gen_dir / schema_tag]
    by_suffix = f" ({by})" if by else ""
    log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
    return new_dirs, log_line


def apply_bump(root: Path, new_dirs: list[Path], log_line: str) -> None:
    for d in new_dirs:
        d.mkdir(parents=True, exist_ok=True)
    versions_path = root / "VERSIONS.md"
    if not versions_path.exists():
        versions_path.write_text("# Dataset versions\n\n")
    with versions_path.open("a") as f:
        f.write(log_line + "\n")


def print_status(root: Path) -> None:
    raw_root = root / "raw"
    if not raw_root.is_dir():
        print(f"no raw/ tree found under {root}")
        return
    for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
        kind = kind_dir.name
        gens = sorted(
            int(m.group(1))
            for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
            if m
        )
        print(f"kind={kind}")
        for gen in gens:
            gen_tag = f"gen{gen}"
            schema_dir = root / "processed" / kind / gen_tag
            schemas = sorted(
                int(m.group(1))
                for m in (
                    SCHEMA_RE.match(p.name)
                    for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
                    if p.is_dir()
                )
                if m
            )
            schema_str = ", ".join(f"schema{s}" for s in schemas) or "(none)"
            print(f"  {gen_tag}: {schema_str}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--root",
        default="/ceph/lbogner/geant_steps",
        help="Dataset root (default: /ceph/lbogner/geant_steps)",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    kind_help = "steps | hits | ... (default: steps)"

    p_gen = sub.add_parser("bump-gen", help="Cut a new raw generation")
    p_gen.add_argument("--kind", default="steps", help=kind_help)
    p_gen.add_argument("--reason", required=True, help="Why this gen exists")
    p_gen.add_argument("--by", default=None, help="Attribution (default: git user.name)")
    p_gen.add_argument("--date", default=None, help="Override date (default: today, ISO)")
    p_gen.add_argument("--execute", action="store_true", help="Apply (default: dry run)")

    p_schema = sub.add_parser("bump-schema", help="Cut a new schema within a gen")
    p_schema.add_argument("--kind", default="steps", help=kind_help)
    p_schema.add_argument("--gen", required=True, help="Existing gen tag, e.g. gen1")
    p_schema.add_argument("--reason", required=True, help="Why this schema exists")
    p_schema.add_argument("--by", default=None, help="Attribution (default: git user.name)")
    p_schema.add_argument("--date", default=None, help="Override date (default: today, ISO)")
    p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)")

    sub.add_parser("status", help="List existing gens/schemas per kind")

    args = parser.parse_args()
    root = Path(args.root)
    if not root.is_dir():
        parser.error(f"{root} is not a directory")

    if args.command == "status":
        print_status(root)
        return

    date = args.date or dt.date.today().isoformat()
    by = args.by if args.by is not None else _git_user_name()

    if args.command == "bump-gen":
        new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date)
    else:
        new_dirs, log_line = plan_bump_schema(
            root, args.kind, args.gen, args.reason, by, date
        )

    print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===")
    print("new directories:")
    for d in new_dirs:
        print(f"  {d}")
    print("VERSIONS.md entry:")
    print(f"  {log_line}")

    if not args.execute:
        print("\nDry run only — pass --execute to apply.")
        return

    apply_bump(root, new_dirs, log_line)
    print("\nDone.")


if __name__ == "__main__":
    main()