diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 5bc5b2a..2187806 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -15,8 +15,10 @@ 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-gen --kind steps --to gen5 --reason "skip to gen5" bump_dataset_version.py bump-schema --kind steps --gen gen1 --reason "added e_sec column" - bump_dataset_version.py update-manifest pools/pbwo4/full.manifest [--schema schema2] + bump_dataset_version.py bump-schema --kind steps --gen gen1 --to schema3 --reason "skip to schema3" + bump_dataset_version.py update-manifest pools/pbwo4/full.manifest [--schema schema2] [--gen gen2] bump_dataset_version.py create-manifest --output pools/pbwo4/train.manifest a.parquet b.parquet bump_dataset_version.py status """ @@ -32,6 +34,14 @@ GEN_RE = re.compile(r"^gen(\d+)$") SCHEMA_RE = re.compile(r"^schema(\d+)$") +def _find_component_idx(abs_path: Path, pattern: re.Pattern) -> int | None: + """Return the index of the first path component matching *pattern*, or None.""" + for i, part in enumerate(abs_path.parts): + if pattern.match(part): + return i + return None + + 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(): @@ -56,18 +66,24 @@ def _git_user_name() -> str | None: def plan_bump_gen( - root: Path, kind: str, reason: str, by: str | None, date: str + root: Path, kind: str, reason: str, by: str | None, date: str, + target: str | None = None, ) -> 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), + """New gen tag is one past the highest seen under raw/ or processed/ for *kind*, + or *target* if explicitly provided.""" + if target is not None: + if not GEN_RE.match(target): + raise SystemExit(f"error: --to must look like 'genN', got {target!r}") + gen_tag = target + else: + next_gen = ( + max( + _max_index(root / "raw" / kind, GEN_RE), + _max_index(root / "processed" / kind, GEN_RE), + ) + + 1 ) - + 1 - ) - gen_tag = f"gen{next_gen}" + gen_tag = f"gen{next_gen}" new_dirs = [ root / "raw" / kind / gen_tag, root / "processed" / kind / gen_tag / "schema1", @@ -78,7 +94,8 @@ def plan_bump_gen( def plan_bump_schema( - root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str + root: Path, kind: str, gen_tag: str, reason: str, by: str | None, date: str, + target: str | None = None, ) -> tuple[list[Path], str]: if not GEN_RE.match(gen_tag): raise SystemExit(f"error: --gen must look like 'genN', got {gen_tag!r}") @@ -88,8 +105,13 @@ def plan_bump_schema( 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}" + if target is not None: + if not SCHEMA_RE.match(target): + raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}") + schema_tag = target + else: + 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}" @@ -139,18 +161,12 @@ def print_status(root: Path) -> None: # update-manifest # --------------------------------------------------------------------------- -def _find_schema_idx(abs_path: Path) -> int | None: - """Return the index of the first schemaN component in abs_path.parts, or None.""" - for i, part in enumerate(abs_path.parts): - if SCHEMA_RE.match(part): - return i - return None - - def plan_update_manifest( - manifest_path: Path, target_schema: str | None + manifest_path: Path, + target_schema: str | None, + target_gen: str | None = None, ) -> tuple[list[tuple[str, str | None]], list[Path]]: - """Parse a manifest and plan schema replacements for each data line. + """Parse a manifest and plan gen/schema replacements for each data line. Returns: lines: list of (original_line, new_relative_path_or_None) @@ -174,30 +190,39 @@ def plan_update_manifest( continue old_abs = (manifest_dir / stripped).resolve() - schema_idx = _find_schema_idx(old_abs) - if schema_idx is None: - result.append((raw, None)) - continue - parts = list(old_abs.parts) - old_schema = parts[schema_idx] - gen_dir = Path(*parts[:schema_idx]) + changed = False - if target_schema is not None: - new_schema = target_schema - else: - if gen_dir not in schema_cache: - n = _max_index(gen_dir, SCHEMA_RE) - if n == 0: - raise SystemExit(f"error: no schema dirs found under {gen_dir}") - schema_cache[gen_dir] = f"schema{n}" - new_schema = schema_cache[gen_dir] + # Apply gen replacement first (shifts subsequent indices). + if target_gen is not None: + gen_idx = _find_component_idx(Path(*parts), GEN_RE) + if gen_idx is not None and parts[gen_idx] != target_gen: + parts[gen_idx] = target_gen + changed = True - if new_schema == old_schema: + schema_idx = _find_component_idx(Path(*parts), SCHEMA_RE) + if schema_idx is not None: + old_schema = parts[schema_idx] + gen_dir = Path(*parts[:schema_idx]) + + if target_schema is not None: + new_schema = target_schema + else: + if gen_dir not in schema_cache: + n = _max_index(gen_dir, SCHEMA_RE) + if n == 0: + raise SystemExit(f"error: no schema dirs found under {gen_dir}") + schema_cache[gen_dir] = f"schema{n}" + new_schema = schema_cache[gen_dir] + + if new_schema != old_schema: + parts[schema_idx] = new_schema + changed = True + + if not changed: result.append((raw, None)) continue - parts[schema_idx] = new_schema new_abs = Path(*parts) if not new_abs.exists(): missing.append(new_abs) @@ -306,6 +331,10 @@ def main() -> None: 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( + "--to", default=None, metavar="genN", + help="Target gen tag (default: one past the current highest)", + ) 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") @@ -314,11 +343,15 @@ def main() -> None: 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( + "--to", default=None, metavar="schemaN", + help="Target schema tag (default: one past the current highest)", + ) p_schema.add_argument("--execute", action="store_true", help="Apply (default: dry run)") p_update = sub.add_parser( "update-manifest", - help="Repoint manifest(s) to a new schema, verifying all target files exist", + help="Repoint manifest(s) to a new gen and/or schema, verifying all target files exist", ) p_update.add_argument( "manifests", nargs="+", metavar="MANIFEST", help="One or more .manifest files to update" @@ -329,6 +362,12 @@ def main() -> None: metavar="schemaN", help="Target schema tag (default: highest schema found in the same gen dir)", ) + p_update.add_argument( + "--gen", + default=None, + metavar="genN", + help="Target gen tag (default: keep existing gen)", + ) p_update.add_argument("--execute", action="store_true", help="Write updated manifests (default: dry run)") p_create = sub.add_parser( @@ -370,9 +409,9 @@ def main() -> None: 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) + new_dirs, log_line = plan_bump_gen(root, args.kind, args.reason, by, date, args.to) else: - new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date) + new_dirs, log_line = plan_bump_schema(root, args.kind, args.gen, args.reason, by, date, args.to) print(f"=== {'EXECUTING' if args.execute else 'DRY RUN'} ===") print("new directories:") @@ -391,13 +430,15 @@ def main() -> None: if args.command == "update-manifest": if args.schema and not SCHEMA_RE.match(args.schema): parser.error(f"--schema must look like 'schemaN', got {args.schema!r}") + if args.gen and not GEN_RE.match(args.gen): + parser.error(f"--gen must look like 'genN', got {args.gen!r}") all_plans: list[tuple[Path, list[tuple[str, str | None]]]] = [] all_missing: list[Path] = [] for raw in args.manifests: mp = Path(raw) - plan, missing = plan_update_manifest(mp, args.schema) + plan, missing = plan_update_manifest(mp, args.schema, args.gen) all_plans.append((mp.resolve(), plan)) all_missing.extend(missing) diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py index 6f3bc30..ccdf3ce 100644 --- a/tests/test_bump_dataset_version.py +++ b/tests/test_bump_dataset_version.py @@ -74,6 +74,40 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path): pass +def test_bump_gen_to_specific_tag(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5") + assert dirs[0] == tmp_path / "raw" / "steps" / "gen5" + assert "`gen5`" in log_line + + +def test_bump_gen_rejects_invalid_to_tag(tmp_path): + try: + plan_bump_gen(tmp_path, "steps", "bad tag", None, "2026-01-01", target="v5") + assert False, "expected SystemExit" + except SystemExit: + pass + + +def test_bump_schema_to_specific_tag(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True) + dirs, log_line = plan_bump_schema( + tmp_path, "steps", "gen1", "jump to schema5", None, "2026-01-01", target="schema5" + ) + assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema5"] + assert "`schema5`" in log_line + + +def test_bump_schema_rejects_invalid_to_tag(tmp_path): + (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True) + try: + plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3") + assert False, "expected SystemExit" + except SystemExit: + pass + + def test_apply_bump_creates_dirs_and_appends_log(tmp_path): dirs, log_line = plan_bump_gen(tmp_path, "steps", "reason A", "alice", "2026-01-01") apply_bump(tmp_path, dirs, log_line) @@ -185,6 +219,40 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path): assert lines[2][1] is not None # the data line was updated +def test_update_manifest_bumps_gen(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n") + + lines, missing = plan_update_manifest(manifest, None, target_gen="gen2") + assert missing == [] + changed = [(old, new) for old, new in lines if new is not None] + assert len(changed) == 1 + assert "gen2" in changed[0][1] + assert "gen1" not in changed[0][1] + + +def test_update_manifest_bumps_gen_and_schema(tmp_path): + parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet" + _make_parquet(parquet) + + manifest_dir = tmp_path / "pools" / "pbwo4" + manifest_dir.mkdir(parents=True) + manifest = manifest_dir / "full.manifest" + manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n") + + lines, missing = plan_update_manifest(manifest, "schema3", target_gen="gen2") + assert missing == [] + changed = [(old, new) for old, new in lines if new is not None] + assert len(changed) == 1 + assert "gen2" in changed[0][1] + assert "schema3" in changed[0][1] + + def test_apply_update_manifest_writes_file(tmp_path): parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet" _make_parquet(parquet)