Add header insertion support to tools/tutorial-cell-updater.py

asumagic · asumagic · commit 0c61d9c16f91 · 2024-09-16T13:16:48.000+02:00
diff --git a/docs/tutorials/notebook-header.md b/docs/tutorials/notebook-header.md
@@ -0,0 +1,6 @@
+<!-- This cell is automatically updated by tools/tutorial-cell-updater.py -->
+<!-- The contents are initialized from tutorials/notebook-header.md -->
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>](https://colab.research.google.com/github/speechbrain/speechbrain/blob/develop/{tutorialpath})
+to execute or view/download this notebook on
+[GitHub](https://github.com/speechbrain/speechbrain/tree/develop/{tutorialpath})
diff --git a/tools/tutorial-cell-updater.py b/tools/tutorial-cell-updater.py
@@ -1,27 +1,58 @@
 #!/usr/bin/env python3
 
-import nbformat
 import glob
-import logging
 import json
+import logging
+
 logging.basicConfig(level=logging.INFO)
 
+with open("tutorials/notebook-header.md", "r", encoding="utf-8") as header_file:
+    header_contents = header_file.read()
+
+
+def find_first_cell_with_tag(cell_list, tag_to_find):
+    for cell in cell_list:
+        tags = cell.get("metadata", {}).get("tags", {})
+        if tag_to_find in tags:
+            return cell
+
+    return None
+
+
+def update_header(header_cell, path):
+    header_cell.update(
+        {
+            "cell_type": "markdown",
+            "metadata": {"id": "sb_auto_header", "tags": ["sb_auto_header"]},
+            "source": header_contents.replace(
+                "{tutorialpath}", path
+            ).splitlines(True),
+        }
+    )
+
+
 def update_notebook(fname):
     logging.info(f"Updating {fname}")
 
+    tutorial_path = fname.replace("./", "")
+
     with open(fname) as f:
         nb = json.load(f)
-        
+
         cells = nb["cells"]
-        
-        for cell in cells:
-            print(list(cell.keys()))
+        header_cell = find_first_cell_with_tag(cells, "sb_auto_header")
+        if header_cell is None:
+            logging.info("Header not found; creating")
+            cells.insert(0, {})
+            header_cell = cells[0]
+
+        update_header(header_cell, tutorial_path)
 
     with open(fname, "w") as wf:
         json.dump(nb, wf, indent=1, ensure_ascii=False)
-        print(file=wf) # print final newline that jupyter adds apparently
+        print(file=wf)  # print final newline that jupyter adds apparently
+
 
 if __name__ == "__main__":
     for fname in glob.glob("./tutorials/**/*.ipynb", recursive=True):
         update_notebook(fname)
-        break
