Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest

bump-gen and bump-schema now accept --to genN/schemaN to target a specific
version instead of always auto-incrementing. update-manifest gains --gen genN
to repoint the gen component of manifest paths (combinable with --schema).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

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*
+    """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."""
+    or *target* if explicitly provided."""
-    next_gen = (
+    if target is not None:
-        max(
+        if not GEN_RE.match(target):
-            _max_index(root / "raw" / kind, GEN_RE),
+            raise SystemExit(f"error: --to must look like 'genN', got {target!r}")
-            _max_index(root / "processed" / kind, GEN_RE),
+        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
+    if target is not None:
-    schema_tag = f"schema{next_schema}"
+        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]
+        changed = False
-        gen_dir = Path(*parts[:schema_idx])
 
-        if target_schema is not None:
+        # Apply gen replacement first (shifts subsequent indices).
-            new_schema = target_schema
+        if target_gen is not None:
-        else:
+            gen_idx = _find_component_idx(Path(*parts), GEN_RE)
-            if gen_dir not in schema_cache:
+            if gen_idx is not None and parts[gen_idx] != target_gen:
-                n = _max_index(gen_dir, SCHEMA_RE)
+                parts[gen_idx] = target_gen
-                if n == 0:
+                changed = True
-                    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:
+        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)
 

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)