From 12082b13c70b0b2373c22a580633a8218def25ae Mon Sep 17 00:00:00 2001
From: Mark McDonald <macd@google.com>
Date: Mon, 29 Jan 2024 14:36:28 +0800
Subject: [PATCH] Add support for `google` in notebook metadata.

These fields will be used to support document-specific metadata that
needs to be pushed to the publishing system.
---
 tools/tensorflow_docs/tools/nbfmt/__main__.py |  9 ++-
 .../tools/nbfmt/nbfmtmain_test.py             | 74 +++++++++++++++++++
 2 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 tools/tensorflow_docs/tools/nbfmt/nbfmtmain_test.py

diff --git a/tools/tensorflow_docs/tools/nbfmt/__main__.py b/tools/tensorflow_docs/tools/nbfmt/__main__.py
index 9426e6fd690..004fd8e9248 100644
--- a/tools/tensorflow_docs/tools/nbfmt/__main__.py
+++ b/tools/tensorflow_docs/tools/nbfmt/__main__.py
@@ -99,16 +99,16 @@ 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 +128,11 @@ 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"])
+  metadata["google"] = google
+
   data["metadata"] = metadata
 
 
diff --git a/tools/tensorflow_docs/tools/nbfmt/nbfmtmain_test.py b/tools/tensorflow_docs/tools/nbfmt/nbfmtmain_test.py
new file mode 100644
index 00000000000..e7a8849b05f
--- /dev/null
+++ b/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()