Allow Notebook output to be preserved

examples/README.adoc

@@ -0,0 +1,15 @@
+# Examples
+
+This folder contains example notebooks on how to use the `graphdatascience` python client.
+
+
+## Custom tooling for notebooks
+
+As for the rest of this repository, you can run `../scripts/makestyle` and `../scripts/checkstyle` to apply and check the style of the notebook.
+By default, `makestyle` will remove all cell outputs. If you want to preserve some outputs, tag the cell with `preserve-output`.
+
+To generate documenatation, you can run `../scripts/nb2doc/convert`.
+
+If you only want to let CI run the notebook given a certain condition, tag a given cell in the notebook with `verify-version`.
+As the name suggests, the tag was introduced to only run for given GDS server versions.
+To make sure certain cells are run even in case of failure, tag the cell with `teardown`.

scripts/checkstyle

@@ -17,19 +17,12 @@ NOTEBOOKS="./examples/*.ipynb" # ./examples/dev/*.ipynb"
 for f in $NOTEBOOKS
 do
   NB=$(cat $f)
-  FORMATTED_NB=$(python -m jupyter nbconvert \
-    --clear-output \
-    --stdout \
-    --ClearOutputPreprocessor.enabled=True \
-    --ClearMetadataPreprocessor.enabled=True \
-    --ClearMetadataPreprocessor.preserve_cell_metadata_mask='tags' \
-    --log-level CRITICAL \
-    $f)
+  FORMATTED_NB=$(python scripts/clean_notebooks.py -i "$f" -o stdout)
 
   if [[ "$FORMATTED_NB" != "$NB" ]];
   then
-    echo "Notebook $f is not correctly formatted"
-    diff --color=always --suppress-common-lines --minimal --side-by-side $NB $FORMATTED_NB
+    echo "Notebook $f is not correctly formatted. See diff above for more details."
+    diff --color=always --suppress-common-lines --minimal --side-by-side <(echo "$NB") <(echo "$FORMATTED_NB")
     exit 1
   fi
 done

scripts/clean_notebooks.py

@@ -0,0 +1,75 @@
+# reasons for not using nbconvert cli tool:
+# * cannot keep output based on a given tag
+
+import argparse
+import logging
+from enum import Enum
+from pathlib import Path
+
+import nbconvert
+from nbconvert.preprocessors import Preprocessor
+
+PRESERVE_CELL_OUTPUT_KEY = "preserve-output"
+METADATA_TAG_KEY = "tags"
+
+
+class OutputMode(Enum):
+    STDOUT = "stdout"
+    INPLACE = "inplace"
+
+
+class CustomClearOutputPreprocessor(Preprocessor):
+    """
+    Removes the output from all code cells in a notebook.
+    Option to keep cell output for cells with a given metadata tag
+    """
+
+    def preprocess_cell(self, cell, resources, cell_index):
+        """
+        Apply a transformation on each cell. See base.py for details.
+        """
+        if cell.cell_type == "code" and PRESERVE_CELL_OUTPUT_KEY not in cell["metadata"].get(METADATA_TAG_KEY, []):
+            cell.outputs = []
+            cell.execution_count = None
+        return cell, resources
+
+
+def main(input_path: Path, output_mode: str) -> None:
+    logger = logging.getLogger("NotebookCleaner")
+    logger.info(f"Cleaning notebooks from `{input_path}`, mode: `{output_mode}`")
+
+    exporter = nbconvert.NotebookExporter()
+
+    metadata_cleaner = nbconvert.preprocessors.ClearMetadataPreprocessor(preserve_cell_metadata_mask=METADATA_TAG_KEY)
+    output_cleaner = CustomClearOutputPreprocessor()
+
+    exporter.register_preprocessor(metadata_cleaner, enabled=True)
+    exporter.register_preprocessor(output_cleaner, enabled=True)
+
+    if input_path.is_file():
+        notebooks = [input_path]
+    else:
+        notebooks = [f for f in input_path.iterdir() if f.is_file() and f.suffix == ".ipynb"]
+
+    logger.info(f"Formatting {len(notebooks)} notebooks.")
+
+    for notebook in notebooks:
+        output = exporter.from_filename(notebook)
+
+        formatted_notebook = output[0]
+
+        if output_mode == OutputMode.INPLACE:
+            with notebook.open(mode="w") as file:
+                file.write(formatted_notebook)
+        elif output_mode == OutputMode.STDOUT:
+            print(formatted_notebook)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--output", choices=[e.value for e in OutputMode])
+    parser.add_argument("-i", "--input", default="examples", help="path to the notebook file or folder")
+
+    args = parser.parse_args()
+
+    main(Path(args.input), OutputMode(args.output))

scripts/makestyle

@@ -13,13 +13,4 @@ if [ "${SKIP_NOTEBOOKS:-false}" == "true" ]; then
   exit 0
 fi
 
-echo "Cleaning notebooks"
-python -m jupyter nbconvert \
---clear-output \
---inplace \
---ClearOutputPreprocessor.enabled=True \
---ClearMetadataPreprocessor.enabled=True \
-  --ClearMetadataPreprocessor.preserve_cell_metadata_mask='tags' \
---log-level CRITICAL \
-./examples/*.ipynb \
-./examples/dev/*.ipynb
+python scripts/clean_notebooks.py -i examples/ -o inplace