From 4b3fb772b8cfda8a57165ba1002213ca73c818e2 Mon Sep 17 00:00:00 2001
From: Umut <umutsahin@protonmail.com>
Date: Mon, 23 Aug 2021 17:00:34 +0300
Subject: [PATCH] chore: integrate checking whether notebooks are sanitized
 into the pcc target

---
 Makefile                                 |  6 ++++-
 script/nbmake_utils/notebook_sanitize.py | 32 ++++++++++++++++--------
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/Makefile b/Makefile
index eca4d4c49..f1ed317a9 100644
--- a/Makefile
+++ b/Makefile
@@ -20,6 +20,10 @@ check_python_format:
 	poetry run env bash ./script/source_format/format_python.sh --dir hdk --dir tests --dir benchmarks --check
 .PHONY: check_python_format
 
+check_strip_nb:
+	poetry run python ./script/nbmake_utils/notebook_sanitize.py examples --check
+.PHONY: strip_nb
+
 pylint:
 	poetry run pylint --rcfile=pylintrc hdk tests benchmarks
 .PHONY: pylint
@@ -38,7 +42,7 @@ pcc:
 	@$(MAKE) --keep-going --jobs $$(nproc) --output-sync --no-print-directory pcc_internal
 .PHONY: pcc
 
-pcc_internal: check_python_format python_linting mypy_ci pydocstyle
+pcc_internal: check_python_format check_strip_nb python_linting mypy_ci pydocstyle
 .PHONY: pcc_internal
 
 pytest:
diff --git a/script/nbmake_utils/notebook_sanitize.py b/script/nbmake_utils/notebook_sanitize.py
index b553f8d52..adcd4499e 100644
--- a/script/nbmake_utils/notebook_sanitize.py
+++ b/script/nbmake_utils/notebook_sanitize.py
@@ -1,21 +1,33 @@
+import argparse
 import json
-import sys
 
 from pathlib import Path
 
 
 def main():
-    path_to_glob = Path(sys.argv[1])
-    notebooks = path_to_glob.glob("*.ipynb")
+    parser = argparse.ArgumentParser(description='Sanitizer for Jupyter Notebooks')
 
-    for notebook_file in notebooks:
-        with open(notebook_file, "r") as f:
-            notebook_dict = json.load(f)
-            notebook_dict["metadata"] = {}
+    parser.add_argument('base', type=str, help='directory which contains the notebooks')
+    parser.add_argument('--check', action='store_true', help='flag to enable just checking mode')
 
-        with open(notebook_file, "w", newline="\n") as f:
-            json.dump(notebook_dict, f, indent=1, ensure_ascii=False)
-            f.write("\n")
+    args = parser.parse_args()
+
+    base = Path(args.base)
+    notebooks = base.glob("*.ipynb")
+
+    for notebook in notebooks:
+        with open(notebook, "r") as f:
+            content = json.load(f)
+
+        if args.check:
+            if len(content["metadata"]) != 0:
+                print("Notebooks are not sanitized. Please run `make conformance`.")
+                exit(1)
+        else:
+            content["metadata"] = {}
+            with open(notebook, "w", newline="\n") as f:
+                json.dump(content, f, indent=1, ensure_ascii=False)
+                f.write("\n")
 
 
 if __name__ == "__main__":