docs: add script for nb stripping, add to docs ci

Author: janfb

.github/workflows/build_docs.yml

@@ -34,6 +34,10 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install .[doc]
 
+    - name: strip output except plots and prints from tutorial notebooks
+      run: |
+        python tests/strip_notebook_outputs.py tutorials/
+
     - name: convert notebooks to markdown
       run: |
         cd docs

pyproject.toml

@@ -51,6 +51,8 @@ doc = [
     # Documentation
     "jupyter_contrib_nbextensions",
     "notebook <= 6.4.12",
+    "nbconvert",
+    "nbformat",
     "traitlets <= 5.9.0",
     "ipython <= 8.9.0",
     "mkdocs",

tests/strip_notebook_outputs.py

@@ -0,0 +1,58 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+# Script to strip outputs of Jupyter notebooks except for plots and prints so that there
+# are displayed nicely on the documentation website.
+
+import os
+
+import nbformat
+from nbconvert.preprocessors import Preprocessor
+
+
+class StripOutputExceptPlotsAndPrintsPreprocessor(Preprocessor):
+    def preprocess_cell(self, cell, resources, cell_index):
+        if cell.cell_type == 'code':
+            # Retain outputs that contain either 'image/png' or 'stdout'
+            cell.outputs = [
+                output
+                for output in cell.outputs
+                if ('data' in output and 'image/png' in output['data'])
+                or (
+                    output.get('output_type') == 'stream'
+                    and output.get('name') == 'stdout'
+                )
+            ]
+        return cell, resources
+
+
+def strip_output_except_plots_and_prints(notebook_path):
+    with open(notebook_path, 'r') as f:
+        nb = nbformat.read(f, as_version=4)
+
+    preprocessor = StripOutputExceptPlotsAndPrintsPreprocessor()
+    nb, _ = preprocessor.preprocess(nb, {})
+
+    with open(notebook_path, 'w') as f:
+        nbformat.write(nb, f)
+
+
+def strip_notebooks_in_directory(directory):
+    """Strips outputs of all notebooks in a directory except for plots and prints."""
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".ipynb"):
+                notebook_path = os.path.join(root, file)
+                strip_output_except_plots_and_prints(notebook_path)
+
+
+# Example usage
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) != 2:
+        print("Usage: python strip_notebook_outputs.py <notebooks_directory>")
+        sys.exit(1)
+
+    notebooks_directory = sys.argv[1]
+    strip_notebooks_in_directory(notebooks_directory)