Merge pull request #2294 from markmcd:nbfmt_devsite

PiperOrigin-RevId: 604399188

tools/tensorflow_docs/tools/nbfmt/__main__.py

@@ -99,16 +99,17 @@ def clean_root(data: Dict[str, Any], filepath: pathlib.Path) -> None:
       data, keep=["cells", "metadata", "nbformat_minor", "nbformat"])
   # All metadata is optional according to spec, but we use some of it.
   notebook_utils.del_entries_except(
-      data["metadata"], keep=["accelerator", "colab", "kernelspec"])
+      data["metadata"], keep=["accelerator", "colab", "kernelspec", "google"]
+  )
 
   metadata = data.get("metadata", {})
-  colab = metadata.get("colab", {})
 
   # Set top-level notebook defaults.
   data["nbformat"] = 4
   data["nbformat_minor"] = 0
 
   # Colab metadata
+  colab = metadata.get("colab", {})
   notebook_utils.del_entries_except(
       colab, keep=["collapsed_sections", "name", "toc_visible"])
   colab["name"] = os.path.basename(filepath)
@@ -128,6 +129,15 @@ def clean_root(data: Dict[str, Any], filepath: pathlib.Path) -> None:
   kernelspec["display_name"] = supported_kernels[kernel_name]
   metadata["kernelspec"] = kernelspec
 
+  # Google metadata
+  google = metadata.get("google", {})
+  notebook_utils.del_entries_except(google, keep=["keywords", "image_path"])
+  # Don't add the field if it's empty.
+  if google:
+    metadata["google"] = google
+  else:
+    metadata.pop("google", None)
+
   data["metadata"] = metadata
 
 

tools/tensorflow_docs/tools/nbfmt/nbfmtmain_test.py

@@ -0,0 +1,74 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for nbfmt."""
+import pathlib
+import unittest
+from nbformat import notebooknode
+from tensorflow_docs.tools.nbfmt import __main__ as nbfmt
+
+
+class NotebookFormatTest(unittest.TestCase):
+
+  def test_metadata_cleansing(self):
+    subject_notebook = notebooknode.NotebookNode({
+        "cells": [],
+        "metadata": {
+            "unknown": ["delete", "me"],
+            "accelerator": "GPU",
+            "colab": {
+                "name": "/this/is/clobbered.ipynb",
+                "collapsed_sections": [],
+                "deleteme": "pls",
+            },
+            "kernelspec": {
+                "display_name": "Python 2 foreverrrr",
+                "name": "python2",
+                "deleteme": "deldeldel",
+            },
+            "google": {
+                "keywords": ["one", "two"],
+                "image_path": "/foo/img.png",
+                "more_stuff": "delete me",
+            },
+        },
+    })
+
+    expected_notebook = notebooknode.NotebookNode({
+        "cells": [],
+        "metadata": {
+            "accelerator": "GPU",
+            "colab": {
+                "name": "test.ipynb",
+                "collapsed_sections": [],
+                "toc_visible": True,
+            },
+            "kernelspec": {
+                "display_name": "Python 3",
+                "name": "python3",
+            },
+            "google": {
+                "keywords": ["one", "two"],
+                "image_path": "/foo/img.png",
+            },
+        },
+        "nbformat": 4,
+        "nbformat_minor": 0,
+    })
+
+    nbfmt.clean_root(subject_notebook, pathlib.Path("/path/test.ipynb"))
+    self.assertEqual(subject_notebook, expected_notebook)
+
+
+if __name__ == "__main__":
+  unittest.main()