Skip to content

Commit 0583008

Browse files
committed
Allow keeping certain output cells
1 parent 657e4a4 commit 0583008

File tree

3 files changed

+77
-18
lines changed

3 files changed

+77
-18
lines changed

scripts/checkstyle

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,7 @@ NOTEBOOKS="./examples/*.ipynb" # ./examples/dev/*.ipynb"
1717
for f in $NOTEBOOKS
1818
do
1919
NB=$(cat $f)
20-
FORMATTED_NB=$(python -m jupyter nbconvert \
21-
--clear-output \
22-
--stdout \
23-
--ClearOutputPreprocessor.enabled=True \
24-
--ClearMetadataPreprocessor.enabled=True \
25-
--ClearMetadataPreprocessor.preserve_cell_metadata_mask='tags' \
26-
--log-level CRITICAL \
27-
$f)
20+
FORMATTED_NB=$(python scripts/clean_notebooks.py -i "$f" -o stdout)
2821

2922
if [[ "$FORMATTED_NB" != "$NB" ]];
3023
then

scripts/clean_notebooks.py

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
# reasons for not using nbconvert cli tool:
2+
# * cannot keep output based on a given tag
3+
4+
import argparse
5+
import logging
6+
from enum import Enum
7+
from pathlib import Path
8+
9+
import nbconvert
10+
from nbconvert.preprocessors import Preprocessor
11+
12+
PRESERVE_CELL_OUTPUT_KEY = "preserve-output"
13+
METADATA_TAG_KEY = "tags"
14+
15+
16+
class OutputMode(Enum):
17+
STDOUT = "stdout"
18+
INPLACE = "inplace"
19+
20+
21+
class CustomClearOutputPreprocessor(Preprocessor):
22+
"""
23+
Removes the output from all code cells in a notebook.
24+
Option to keep cell output for cells with a given metadata tag
25+
"""
26+
27+
def preprocess_cell(self, cell, resources, cell_index):
28+
"""
29+
Apply a transformation on each cell. See base.py for details.
30+
"""
31+
if cell.cell_type == "code" and PRESERVE_CELL_OUTPUT_KEY not in cell["metadata"].get(METADATA_TAG_KEY, []):
32+
cell.outputs = []
33+
cell.execution_count = None
34+
return cell, resources
35+
36+
37+
def main(input_path: Path, output_mode: str) -> None:
38+
logger = logging.getLogger("NotebookCleaner")
39+
logger.info(f"Cleaning notebooks from `{input_path}`, mode: `{output_mode}`")
40+
41+
exporter = nbconvert.NotebookExporter()
42+
43+
metadata_cleaner = nbconvert.preprocessors.ClearMetadataPreprocessor(preserve_cell_metadata_mask=METADATA_TAG_KEY)
44+
output_cleaner = CustomClearOutputPreprocessor()
45+
46+
exporter.register_preprocessor(metadata_cleaner, enabled=True)
47+
exporter.register_preprocessor(output_cleaner, enabled=True)
48+
49+
if input_path.is_file():
50+
notebooks = [input_path]
51+
else:
52+
notebooks = [f for f in input_path.iterdir() if f.is_file() and f.suffix == ".ipynb"]
53+
54+
logger.info(f"Formatting {len(notebooks)} notebooks.")
55+
56+
for notebook in notebooks:
57+
output = exporter.from_filename(notebook)
58+
59+
formatted_notebook = output[0]
60+
61+
if output_mode == OutputMode.INPLACE:
62+
with notebook.open(mode="w") as file:
63+
file.write(formatted_notebook)
64+
elif output_mode == OutputMode.STDOUT:
65+
print(formatted_notebook)
66+
67+
68+
if __name__ == "__main__":
69+
parser = argparse.ArgumentParser()
70+
parser.add_argument("-o", "--output", choices=[e.value for e in OutputMode])
71+
parser.add_argument("-i", "--input", default="examples", help="path to the notebook file or folder")
72+
73+
args = parser.parse_args()
74+
75+
main(Path(args.input), OutputMode(args.output))

scripts/makestyle

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,4 @@ if [ "${SKIP_NOTEBOOKS:-false}" == "true" ]; then
1313
exit 0
1414
fi
1515

16-
echo "Cleaning notebooks"
17-
python -m jupyter nbconvert \
18-
--clear-output \
19-
--inplace \
20-
--ClearOutputPreprocessor.enabled=True \
21-
--ClearMetadataPreprocessor.enabled=True \
22-
--ClearMetadataPreprocessor.preserve_cell_metadata_mask='tags' \
23-
--log-level CRITICAL \
24-
./examples/*.ipynb \
25-
./examples/dev/*.ipynb
16+
python scripts/clean_notebooks.py -i examples/ -o inplace

0 commit comments

Comments
 (0)