wehub-resource-sync
119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
"""Add missing Papermill parameters cells to notebooks.
|
|
|
|
Notebooks that already have a `# %% tags=["parameters"]` cell are skipped.
|
|
For notebooks without one, this inserts an empty parameters cell after the
|
|
first code cell (or after imports if detectable).
|
|
|
|
Usage:
|
|
uv run python tests/add_missing_parameters_cells.py [--dry-run]
|
|
"""
|
|
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).parent.parent
|
|
DRY_RUN = "--dry-run" in sys.argv
|
|
|
|
# Parameters cell template
|
|
PARAMS_CELL = """
|
|
# %% tags=["parameters"]
|
|
# Production defaults — Papermill injects overrides for CI
|
|
"""
|
|
|
|
|
|
def find_notebooks() -> list[Path]:
|
|
"""Find all Jupytext notebooks missing a parameters cell."""
|
|
notebooks = []
|
|
for d in sorted(REPO_ROOT.glob("[0-9]*_*")):
|
|
notebooks.extend(sorted(d.glob("**/*.py")))
|
|
for d in sorted((REPO_ROOT / "case_studies").glob("*")):
|
|
if d.is_dir() and d.name not in ("utils", "__pycache__"):
|
|
notebooks.extend(sorted(d.glob("**/*.py")))
|
|
|
|
missing = []
|
|
for nb in notebooks:
|
|
content = nb.read_text()
|
|
if "# %%" not in content:
|
|
continue
|
|
if 'tags=["parameters"]' in content or "tags=['parameters']" in content:
|
|
continue
|
|
missing.append(nb)
|
|
return missing
|
|
|
|
|
|
def add_parameters_cell(nb_path: Path) -> bool:
|
|
"""Insert a parameters cell after the first code cell.
|
|
|
|
Strategy: find the first `# %%` code cell (not markdown, not the header)
|
|
and insert the parameters cell after it. If the first code cell has imports,
|
|
the parameters cell goes after the import block.
|
|
"""
|
|
content = nb_path.read_text()
|
|
lines = content.split("\n")
|
|
|
|
# Find insertion point: after the first code cell
|
|
# Skip: header (---/jupyter metadata), markdown cells
|
|
in_header = False
|
|
found_first_code = False
|
|
insert_after = None
|
|
|
|
for i, line in enumerate(lines):
|
|
stripped = line.strip()
|
|
|
|
# Track Jupytext header
|
|
if i == 0 and stripped == "# ---":
|
|
in_header = True
|
|
continue
|
|
if in_header:
|
|
if stripped == "# ---":
|
|
in_header = False
|
|
continue
|
|
|
|
# Found a code cell
|
|
if stripped == "# %%" and not found_first_code:
|
|
found_first_code = True
|
|
# Look ahead to find end of this cell (next # %% or end)
|
|
for j in range(i + 1, len(lines)):
|
|
if lines[j].strip().startswith("# %%"):
|
|
insert_after = j
|
|
break
|
|
else:
|
|
insert_after = len(lines)
|
|
break
|
|
|
|
if insert_after is None:
|
|
return False
|
|
|
|
# Insert the parameters cell
|
|
new_lines = (
|
|
lines[:insert_after] + PARAMS_CELL.rstrip().split("\n") + [""] + lines[insert_after:]
|
|
)
|
|
new_content = "\n".join(new_lines)
|
|
|
|
if not DRY_RUN:
|
|
nb_path.write_text(new_content)
|
|
return True
|
|
|
|
|
|
def main():
|
|
missing = find_notebooks()
|
|
print(f"Found {len(missing)} notebooks missing parameters cell")
|
|
if DRY_RUN:
|
|
print("(dry run — no changes)")
|
|
|
|
modified = 0
|
|
for nb in missing:
|
|
rel = nb.relative_to(REPO_ROOT)
|
|
ok = add_parameters_cell(nb)
|
|
status = "ADDED" if ok else "SKIPPED"
|
|
if ok:
|
|
modified += 1
|
|
print(f" {status}: {rel}")
|
|
|
|
print(f"\nModified: {modified}/{len(missing)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|