vondp/TorchCode

1#!/usr/bin/env python32"""Add 'Open in Colab' badges to all template notebooks."""3 4import json5from pathlib import Path6 7REPO = "duoan/TorchCode"8BRANCH = "master"9TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates"10BADGE_IMG = "https://colab.research.google.com/assets/colab-badge.svg"11 12 13def colab_url(filename: str) -> str:14    return (15        f"https://colab.research.google.com/github/{REPO}"16        f"/blob/{BRANCH}/templates/{filename}"17    )18 19 20def badge_markdown(filename: str) -> str:21    return f"[![Open In Colab]({BADGE_IMG})]({colab_url(filename)})"22 23 24def process_notebook(path: Path) -> bool:25    with open(path, "r", encoding="utf-8") as f:26        nb = json.load(f)27 28    cells = nb.get("cells", [])29    if not cells or cells[0].get("cell_type") != "markdown":30        return False31 32    source_lines = cells[0]["source"]33    flat = "".join(source_lines) if isinstance(source_lines, list) else source_lines34    if "colab-badge.svg" in flat:35        return False36 37    badge = badge_markdown(path.name)38    cells[0]["source"] = [badge + "\n\n"] + (39        source_lines if isinstance(source_lines, list) else [source_lines]40    )41 42    with open(path, "w", encoding="utf-8") as f:43        json.dump(nb, f, ensure_ascii=False, indent=1)44        f.write("\n")45 46    return True47 48 49def main() -> None:50    updated = 051    for nb_path in sorted(TEMPLATES_DIR.glob("*.ipynb")):52        if process_notebook(nb_path):53            print(f"  ✅ {nb_path.name}")54            updated += 155        else:56            print(f"  ⏭️  {nb_path.name} (already has badge or skipped)")57    print(f"\nDone — updated {updated} notebooks.")58 59 60if __name__ == "__main__":61    main()62