Move the page headers to the templates.

PiperOrigin-RevId: 421556385

Author: MarkDaoust

tools/tensorflow_docs/api_generator/generate_lib.py

@@ -437,27 +437,6 @@ def generate(self) -> Dict[str, Any]:
     return {'toc': toc}
 
 
-def _get_headers(page_info: base_page.PageInfo,
-                 search_hints: bool) -> List[str]:
-  """Returns the list of header lines for this page."""
-  hidden = doc_controls.should_hide_from_search(page_info.py_object)
-  brief_no_backticks = page_info.doc.brief.replace('`', '').strip()
-  headers = []
-  if brief_no_backticks:
-    headers.append(f'description: {brief_no_backticks}')
-
-  # It's easier to read if there's a blank line after the `name:value` headers.
-  if search_hints and not hidden:
-    if headers:
-      headers.append('')
-    headers.append(page_info.get_metadata_html())
-  else:
-    headers.append('robots: noindex')
-    headers.append('')
-
-  return headers
-
-
 def write_docs(
     *,
     output_dir: Union[str, pathlib.Path],
@@ -566,9 +545,14 @@ def write_docs(
 
     # Generate docs for `py_object`, resolving references.
     try:
-      page_info = docs_for_object.docs_for_object(full_name, py_object,
-                                                  parser_config, extra_docs,
-                                                  page_builder_classes)
+      page_info = docs_for_object.docs_for_object(
+          full_name=full_name,
+          py_object=py_object,
+          parser_config=parser_config,
+          extra_docs=extra_docs,
+          search_hints=search_hints,
+          page_builder_classes=page_builder_classes)
+
       if gen_report and not full_name.startswith(
           ('tf.compat.v', 'tf.keras.backend', 'tf.numpy',
            'tf.experimental.numpy')):
@@ -579,9 +563,8 @@ def write_docs(
 
     path = output_dir / parser.documentation_path(full_name)
 
-    content = _get_headers(page_info, search_hints)
-    content.append(page_info.build())
-    text = '\n'.join(content)
+    text = page_info.build()
+
     try:
       path.parent.mkdir(exist_ok=True, parents=True)
       path.write_text(text, encoding='utf-8')

tools/tensorflow_docs/api_generator/generate_lib_test.py

@@ -187,60 +187,6 @@ def test_write(self):
     # Make sure that duplicates are not written
     self.assertTrue((output_dir / 'tf/TestModule/test_function.md').exists())
 
-  def _get_test_page_info(self):
-    page_info = function_page.FunctionPageInfo(
-        full_name='abc', py_object=test_function)
-    docstring_info = parser.DocstringInfo(
-        brief='hello `tensorflow`',
-        docstring_parts=['line1', 'line2'],
-        compatibility={})
-    page_info.set_doc(docstring_info)
-    return page_info
-
-  def test_get_headers_global_hints(self):
-    page_info = self._get_test_page_info()
-    result = '\n'.join(generate_lib._get_headers(page_info, search_hints=True))
-
-    expected = textwrap.dedent("""\
-      description: hello tensorflow
-
-      <div itemscope itemtype="http://developers.google.com/ReferenceObject">
-      <meta itemprop="name" content="abc" />
-      <meta itemprop="path" content="Stable" />
-      </div>
-      """)
-
-    self.assertEqual(expected, result)
-
-  def test_get_headers_global_no_hints(self):
-    page_info = self._get_test_page_info()
-    result = '\n'.join(generate_lib._get_headers(page_info, search_hints=False))
-
-    expected = textwrap.dedent("""\
-      description: hello tensorflow
-      robots: noindex
-      """)
-
-    self.assertEqual(expected, result)
-
-  def test_get_headers_local_no_hints(self):
-    page_info = self._get_test_page_info()
-
-    @doc_controls.hide_from_search
-    def py_object():
-      pass
-
-    page_info.py_object = py_object
-
-    result = '\n'.join(generate_lib._get_headers(page_info, search_hints=True))
-
-    expected = textwrap.dedent("""\
-      description: hello tensorflow
-      robots: noindex
-      """)
-
-    self.assertEqual(expected, result)
-
 
 if __name__ == '__main__':
   absltest.main()

tools/tensorflow_docs/api_generator/pretty_docs/base_page.py

@@ -19,6 +19,7 @@
 from typing import Any, Callable, ClassVar, Dict, List, NamedTuple, Optional, Sequence, Tuple, Type
 
 from tensorflow_docs.api_generator import config
+from tensorflow_docs.api_generator import doc_controls
 from tensorflow_docs.api_generator import parser
 from tensorflow_docs.api_generator import signature as signature_lib
 
@@ -63,6 +64,25 @@ def bottom_compat(self):
   def format_docstring_part(self, part):
     return str(part)
 
+  def get_devsite_headers(self):
+    """Returns the list of header lines for this page."""
+    hidden = doc_controls.should_hide_from_search(self.page_info.py_object)
+    brief_no_backticks = self.page_info.doc.brief.replace('`', '').strip()
+    headers = []
+    if brief_no_backticks:
+      headers.append(f'description: {brief_no_backticks}')
+
+    if self.page_info.search_hints and not hidden:
+      if headers:
+        headers.append('')
+      headers.append(self.page_info.get_metadata_html())
+    else:
+      headers.append('robots: noindex')
+      headers.append('')
+
+    result = '\n'.join(headers)
+    return result
+
 
 class PageInfo:
   """Base-class for api_pages objects.
@@ -85,6 +105,7 @@ def __init__(
       full_name: str,
       py_object: Any,
       extra_docs: Optional[Dict[int, str]] = None,
+      search_hints: bool = True,
   ):
     """Initialize a PageInfo.
 
@@ -97,6 +118,7 @@ def __init__(
     self.full_name = full_name
     self.py_object = py_object
     self._extra_docs = extra_docs
+    self.search_hints = search_hints
 
     self._defined_in = None
     self._aliases = None

tools/tensorflow_docs/api_generator/pretty_docs/docs_for_object.py

@@ -40,10 +40,12 @@
 
 
 def docs_for_object(
+    *,
     full_name: str,
     py_object: Any,
     parser_config: config.ParserConfig,
     extra_docs: Optional[Dict[int, str]] = None,
+    search_hints: bool = True,
     page_builder_classes: Optional[PageBuilderDict] = None,
 ) -> base_page.PageInfo:
   """Return a PageInfo object describing a given object from the TF API.
@@ -85,7 +87,10 @@ def docs_for_object(
     page_info_class = page_builder_classes[obj_type]
 
   page_info = page_info_class(
-      full_name=main_name, py_object=py_object, extra_docs=extra_docs)
+      full_name=main_name,
+      py_object=py_object,
+      search_hints=search_hints,
+      extra_docs=extra_docs)
 
   relative_path = os.path.relpath(
       path='.',

tools/tensorflow_docs/api_generator/pretty_docs/pretty_docs_test.py

@@ -19,8 +19,10 @@
 
 from absl.testing import absltest
 
+from tensorflow_docs.api_generator import doc_controls
 from tensorflow_docs.api_generator import parser
 from tensorflow_docs.api_generator.pretty_docs import base_page
+from tensorflow_docs.api_generator.pretty_docs import function_page
 
 
 class ParserTest(absltest.TestCase):
@@ -70,6 +72,66 @@ def test_no_source_link(self):
        """)
     self.assertEqual(expected, table)
 
+  def _get_test_page_builder(self, search_hints):
+
+    def test_function():
+      pass
+
+    page_info = function_page.FunctionPageInfo(
+        full_name='abc', py_object=test_function, search_hints=search_hints)
+    docstring_info = parser.DocstringInfo(
+        brief='hello `tensorflow`',
+        docstring_parts=['line1', 'line2'],
+        compatibility={})
+    page_info.set_doc(docstring_info)
+    page_builder = function_page.FunctionPageBuilder(page_info)
+    return page_builder
+
+  def test_get_headers_global_hints(self):
+    page_builder = self._get_test_page_builder(search_hints=True)
+    result = page_builder.get_devsite_headers()
+
+    expected = textwrap.dedent("""\
+      description: hello tensorflow
+
+      <div itemscope itemtype="http://developers.google.com/ReferenceObject">
+      <meta itemprop="name" content="abc" />
+      <meta itemprop="path" content="Stable" />
+      </div>
+      """)
+
+    self.assertEqual(expected, result)
+
+  def test_get_headers_global_no_hints(self):
+    page_builder = self._get_test_page_builder(search_hints=False)
+    result = page_builder.get_devsite_headers()
+
+    expected = textwrap.dedent("""\
+      description: hello tensorflow
+      robots: noindex
+      """)
+
+    self.assertEqual(expected, result)
+
+  def test_get_headers_local_no_hints(self):
+    page_builder = self._get_test_page_builder(search_hints=True)
+    result = page_builder.get_devsite_headers()
+
+    @doc_controls.hide_from_search
+    def py_object():
+      pass
+
+    page_builder.page_info.py_object = py_object
+
+    result = page_builder.get_devsite_headers()
+
+    expected = textwrap.dedent("""\
+      description: hello tensorflow
+      robots: noindex
+      """)
+
+    self.assertEqual(expected, result)
+
 
 if __name__ == "__main__":
   absltest.main()

tools/tensorflow_docs/api_generator/pretty_docs/templates/module.jinja

@@ -3,6 +3,7 @@
 {#----------------------------------------------------------------------------#}
 
 {% block header %}
+{{ builder.get_devsite_headers() }}
 # Module: {{page_info.full_name}}
 
 <!-- Insert buttons and diff -->

tools/tensorflow_docs/api_generator/pretty_docs/templates/page.jinja

@@ -11,6 +11,7 @@
 {#----------------------------------------------------------------------------#}
 
 {% block header %}
+{{ builder.get_devsite_headers() }}
 # {{page_info.full_name}}
 
 <!-- Insert buttons and diff -->