ENH: check multiple autolink-concat in set-nb-cells

src/compwa_policy/set_nb_cells.py

@@ -1,12 +1,12 @@
 """Add or update standard cells in a Jupyter notebook.
 
 Notebook servers like Google Colaboratory and Deepnote do not install a package
-automatically, so this has to be done through a code cell. At the same time,
-this cell needs to be hidden from the documentation pages, when viewing through
-Jupyter Lab (Binder), and when viewing as Jupyter slides.
+automatically, so this has to be done through a code cell. At the same time, this cell
+needs to be hidden from the documentation pages, when viewing through Jupyter Lab
+(Binder), and when viewing as Jupyter slides.
 
-This scripts sets the IPython InlineBackend.figure_formats option to SVG. This
-is because the Sphinx configuration can't set this externally.
+This scripts sets the IPython InlineBackend.figure_formats option to SVG. This is
+because the Sphinx configuration can't set this externally.
 
 Notebooks can be ignored by making the first cell a `Markdown cell
 <https://jupyter-notebook.readthedocs.io/en/latest/examples/Notebook/Working%20With%20Markdown%20Cells.html>`_
@@ -24,9 +24,10 @@
 import sys
 from functools import lru_cache
 from textwrap import dedent
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import nbformat
+from nbformat import NotebookNode
 
 from compwa_policy.utilities.notebook import load_notebook
 from compwa_policy.utilities.pyproject import Pyproject
@@ -52,6 +53,11 @@
     "tags": ["remove-cell", "skip-execution"],
     # https://github.com/executablebooks/jupyter-book/issues/833
 }
+__AUTOLINK_CONCAT = """
+```{autolink-concat}
+
+```
+""".strip()
 
 
 def main(argv: Sequence[str] | None = None) -> int:
@@ -86,8 +92,11 @@ def main(argv: Sequence[str] | None = None) -> int:
     )
     args = parser.parse_args(argv)
 
+    failed = False
     for filename in args.filenames:
         cell_id = 0
+        updated = False
+        notebook = load_notebook(filename)
         if args.add_install_cell:
             cell_content = __get_install_cell().strip("\n")
             if args.extras_require:
@@ -96,23 +105,40 @@ def main(argv: Sequence[str] | None = None) -> int:
             if args.additional_packages:
                 packages = [s.strip() for s in args.additional_packages.split(",")]
                 cell_content += " " + " ".join(packages)
-            _update_cell(
-                filename,
+            updated |= _update_cell(
+                notebook,
                 new_content=cell_content,
                 new_metadata=__INSTALL_CELL_METADATA,
                 cell_id=cell_id,
             )
             cell_id += 1
         if args.config_cell:
             config_cell_content = __CONFIG_CELL_CONTENT
-            _update_cell(
+            updated |= _update_cell(
                 filename,
                 new_content=config_cell_content.strip("\n"),
                 new_metadata=__CONFIG_CELL_METADATA,
                 cell_id=cell_id,
             )
-        if args.autolink_concat:
-            _insert_autolink_concat(filename)
+        if args.autolink_concat and not _skip_notebook(
+            notebook, ignore_comment="<!-- no autolink-concat -->"
+        ):
+            updated |= _format_autolink_concat(notebook)
+            updated |= _insert_autolink_concat(notebook)
+            if (n_autolink := _count_autolink_concat(notebook)) > 1:
+                failed |= True
+                print(  # noqa: T201
+                    f"Found {n_autolink} autolink-concat cells in {filename}, should be"
+                    " only one. Please remove duplicates.",
+                    file=sys.stderr,
+                )
+        if updated:
+            print(f"Updated {filename}", file=sys.stderr)  # noqa: T201
+            nbformat.validate(notebook)
+            nbformat.write(notebook, filename)
+            failed |= True
+    if failed:
+        return 1
     return 0
 
 
@@ -127,14 +153,13 @@ def __get_install_cell() -> str:
 
 
 def _update_cell(
-    filename: str,
+    notebook: NotebookNode,
     new_content: str,
     new_metadata: dict,
     cell_id: int,
-) -> None:
-    if _skip_notebook(filename):
-        return
-    notebook = load_notebook(filename)
+) -> bool:
+    if _skip_notebook(notebook, ignore_comment="<!-- no-set-nb-cells -->"):
+        return False
     exiting_cell = notebook["cells"][cell_id]
     new_cell = nbformat.v4.new_code_cell(
         new_content,
@@ -145,42 +170,75 @@ def _update_cell(
         notebook["cells"][cell_id] = new_cell
     else:
         notebook["cells"].insert(cell_id, new_cell)
-    nbformat.validate(notebook)
-    nbformat.write(notebook, filename)
+    return True
+
+
+def _format_autolink_concat(notebook: NotebookNode) -> bool:
+    candidates = [
+        dedent("""
+        ```{autolink-concat}
+        ```
+        """).strip(),
+        dedent("""
+        :::{autolink-concat}
+        :::
+        """).strip(),
+    ]
+    updated = False
+    for cell in notebook["cells"]:
+        if cell["cell_type"] != "markdown":
+            continue
+        for pattern in candidates:
+            source = cast("str", cell["source"])
+            if pattern in source:
+                cell["source"] = source.replace(pattern, __AUTOLINK_CONCAT)
+                updated |= True
+    return updated
 
 
-def _insert_autolink_concat(filename: str) -> None:
-    if _skip_notebook(filename, ignore_statement="<!-- no autolink-concat -->"):
-        return
-    notebook = load_notebook(filename)
-    expected_cell_content = """
+def _insert_autolink_concat(notebook: NotebookNode) -> bool:
+    expected_cell_content = dedent("""
     ```{autolink-concat}
+
     ```
-    """
-    expected_cell_content = dedent(expected_cell_content).strip()
+    """).strip()
+    if any(
+        expected_cell_content in cell["source"]
+        for cell in notebook["cells"]
+        if cell["cell_type"] == "markdown"
+    ):
+        return False
     for cell_id, cell in enumerate(notebook["cells"]):
         if cell["cell_type"] != "markdown":
             continue
-        cell_content: str = cell["source"]
-        if expected_cell_content in cell_content:
-            return
         new_cell = nbformat.v4.new_markdown_cell(expected_cell_content)
         del new_cell["id"]  # following nbformat_minor = 4
         notebook["cells"].insert(cell_id, new_cell)
-        nbformat.validate(notebook)
-        nbformat.write(notebook, filename)
-        return
+        return True
+    return False
 
 
-def _skip_notebook(
-    filename: str, ignore_statement: str = "<!-- no-set-nb-cells -->"
-) -> bool:
-    notebook = load_notebook(filename)
+def _count_autolink_concat(notebook: NotebookNode) -> int:
+    search_terms = [
+        "```{autolink-concat}",
+        ":::{autolink-concat}",
+    ]
+    count = 0
+    for cell in notebook["cells"]:
+        if cell["cell_type"] != "markdown":
+            continue
+        cell_content: str = cell["source"]
+        if any(term in cell_content for term in search_terms):
+            count += 1
+    return count
+
+
+def _skip_notebook(notebook: NotebookNode, ignore_comment: str) -> bool:
     for cell in notebook["cells"]:
         if cell["cell_type"] != "markdown":
             continue
         cell_content: str = cell["source"]
-        if ignore_statement in cell_content.lower():
+        if ignore_comment in cell_content.lower():
             return True
     return False