Add injection sites for customizing page construction.

* For single-page changes, tag an object using `doc_controls.set_custom_page_builder_c lass(obj, cls)
* For global changes pass a dict of `{ObjectType:Type[PageInfo]}` to the DocGenerator's `page_builder_classes` argument.
* Switch generate2 to use the new customization pathways.

PiperOrigin-RevId: 420801142

Author: MarkDaoust

tools/tensorflow_docs/api_generator/doc_controls.py

@@ -73,17 +73,17 @@ def should_hide_from_search(obj) -> bool:
   return hasattr(obj, _NO_SEARCH_HINTS)
 
 
-_CUSTOM_PAGE_CONTENT = "_tf_docs_custom_page_content"
+_CUSTOM_PAGE_BUILDER_CLS = "_tf_docs_custom_page_builder_cls"
 
 
-def set_custom_page_content(obj, content):
+def set_custom_page_builder_cls(obj, cls):
   """Replace most of the generated page with custom content."""
-  setattr(obj, _CUSTOM_PAGE_CONTENT, content)
+  setattr(obj, _CUSTOM_PAGE_BUILDER_CLS, cls)
 
 
-def get_custom_page_content(obj):
+def get_custom_page_builder_cls(obj):
   """Gets custom page content if available."""
-  return getattr(obj, _CUSTOM_PAGE_CONTENT, None)
+  return getattr(obj, _CUSTOM_PAGE_BUILDER_CLS, None)
 
 
 _DO_NOT_DOC = "_tf_docs_do_not_document"

tools/tensorflow_docs/api_generator/generate_lib.py

@@ -470,6 +470,7 @@ def write_docs(
     gen_redirects: bool = True,
     gen_report: bool = True,
     extra_docs: Optional[Dict[int, str]] = None,
+    page_builder_classes: Optional[docs_for_object.PageBuilderDict] = None,
 ):
   """Write previously extracted docs to disk.
 
@@ -498,6 +499,8 @@ def write_docs(
     extra_docs: To add docs for a particular object instance set it's __doc__
       attribute. For some classes (list, tuple, etc) __doc__ is not writable.
       Pass those docs like: `extra_docs={id(obj): "docs"}`
+    page_builder_classes: A optional dict of `{ObjectType:Type[PageInfo]}` for
+        overriding the default page builder classes.
 
   Raises:
     ValueError: if `output_dir` is not an absolute path
@@ -564,7 +567,8 @@ 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)
+                                                  parser_config, extra_docs,
+                                                  page_builder_classes)
       if gen_report and not full_name.startswith(
           ('tf.compat.v', 'tf.keras.backend', 'tf.numpy',
            'tf.experimental.numpy')):
@@ -731,6 +735,7 @@ def __init__(
       gen_redirects: bool = True,
       gen_report: bool = True,
       extra_docs: Optional[Dict[int, str]] = None,
+      page_builder_classes: Optional[docs_for_object.PageBuilderDict] = None,
   ):
     """Creates a doc-generator.
 
@@ -766,6 +771,8 @@ def __init__(
       extra_docs: To add docs for a particular object instance set it's __doc__
         attribute. For some classes (list, tuple, etc) __doc__ is not writable.
         Pass those docs like: `extra_docs={id(obj): "docs"}`
+      page_builder_classes: An optional dict of `{ObjectType:Type[PageInfo]}`
+        for overriding the default page builder classes.
     """
     self._root_title = root_title
     self._py_modules = py_modules
@@ -807,6 +814,7 @@ def __init__(
     self._gen_redirects = gen_redirects
     self._gen_report = gen_report
     self._extra_docs = extra_docs
+    self._page_builder_classes = page_builder_classes
 
   def make_reference_resolver(self, visitor):
     return reference_resolver_lib.ReferenceResolver.from_visitor(
@@ -868,6 +876,7 @@ def build(self, output_dir):
         gen_redirects=self._gen_redirects,
         gen_report=self._gen_report,
         extra_docs=self._extra_docs,
+        page_builder_classes=self._page_builder_classes,
     )
 
     if self.api_cache:

tools/tensorflow_docs/api_generator/pretty_docs/base_page.py

@@ -18,6 +18,7 @@
 import textwrap
 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 parser
 from tensorflow_docs.api_generator import signature as signature_lib
 
@@ -101,8 +102,13 @@ def __init__(
     self._aliases = None
     self._doc = None
 
+  def collect_docs(self, parser_config: config.ParserConfig):
+    """Collects additional information from the `config.ParserConfig`."""
+    pass
+
   def build(self) -> str:
-    """Builds the documentation"""
+    """Builds the documentation."""
+
     cls = self.DEFAULT_BUILDER_CLASS
     return cls(self).build()
 

tools/tensorflow_docs/api_generator/pretty_docs/docs_for_object.py

@@ -17,9 +17,10 @@
 import os
 import posixpath
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Type
 
 from tensorflow_docs.api_generator import config
+from tensorflow_docs.api_generator import doc_controls
 from tensorflow_docs.api_generator import obj_type as obj_type_lib
 from tensorflow_docs.api_generator import parser
 from tensorflow_docs.api_generator.pretty_docs import base_page
@@ -28,12 +29,22 @@
 from tensorflow_docs.api_generator.pretty_docs import module_page
 from tensorflow_docs.api_generator.pretty_docs import type_alias_page
 
+_DEFAULT_PAGE_BUILDER_CLASSES = {
+    obj_type_lib.ObjType.CLASS: class_page.ClassPageInfo,
+    obj_type_lib.ObjType.CALLABLE: function_page.FunctionPageInfo,
+    obj_type_lib.ObjType.MODULE: module_page.ModulePageInfo,
+    obj_type_lib.ObjType.TYPE_ALIAS: type_alias_page.TypeAliasPageInfo,
+}
+
+PageBuilderDict = Dict[obj_type_lib.ObjType, Type[base_page.PageInfo]]
+
 
 def docs_for_object(
     full_name: str,
     py_object: Any,
     parser_config: config.ParserConfig,
     extra_docs: Optional[Dict[int, str]] = None,
+    page_builder_classes: Optional[PageBuilderDict] = None,
 ) -> base_page.PageInfo:
   """Return a PageInfo object describing a given object from the TF API.
 
@@ -47,6 +58,8 @@ def docs_for_object(
     parser_config: A `config.ParserConfig` object.
     extra_docs: Extra docs for symbols like public constants(list, tuple, etc)
       that need to be added to the markdown pages created.
+    page_builder_classes: An optional dict of `{ObjectType:Type[PageInfo]}` for
+        overriding the default page builder classes.
 
   Returns:
     Either a subclass of `pretty_docs.base_page.PageInfo` depending on the type
@@ -63,21 +76,16 @@ def docs_for_object(
   if main_name in duplicate_names:
     duplicate_names.remove(main_name)
 
-  obj_type = obj_type_lib.ObjType.get(py_object)
-  if obj_type is obj_type_lib.ObjType.CLASS:
-    page_info = class_page.ClassPageInfo(
-        full_name=main_name, py_object=py_object, extra_docs=extra_docs)
-  elif obj_type is obj_type_lib.ObjType.CALLABLE:
-    page_info = function_page.FunctionPageInfo(
-        full_name=main_name, py_object=py_object, extra_docs=extra_docs)
-  elif obj_type is obj_type_lib.ObjType.MODULE:
-    page_info = module_page.ModulePageInfo(
-        full_name=main_name, py_object=py_object, extra_docs=extra_docs)
-  elif obj_type is obj_type_lib.ObjType.TYPE_ALIAS:
-    page_info = type_alias_page.TypeAliasPageInfo(
-        full_name=main_name, py_object=py_object, extra_docs=extra_docs)
-  else:
-    raise RuntimeError('Cannot make docs for object {full_name}: {py_object!r}')
+  if page_builder_classes is None:
+    page_builder_classes = _DEFAULT_PAGE_BUILDER_CLASSES
+
+  page_info_class = doc_controls.get_custom_page_builder_cls(py_object)
+  if page_info_class is None:
+    obj_type = obj_type_lib.ObjType.get(py_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)
 
   relative_path = os.path.relpath(
       path='.',

tools/tensorflow_docs/api_generator/pretty_docs/module_page.py

@@ -50,11 +50,6 @@ def build(self) -> str:
 
     parts.append('\n\n')
 
-    custom_content = doc_controls.get_custom_page_content(page_info.py_object)
-    if custom_content is not None:
-      parts.append(custom_content)
-      return ''.join(parts)
-
     if page_info.modules:
       parts.append('## Modules\n\n')
       parts.extend(