Simplify in generate_lib with the ApiTree and Toc classes.

- config.py: Attach the api_tree to the parser_config.
- doc_generator_visitor.py: Use the api_tree to generate the toc.
- toc.py: move root's submodules from children to peers.

PiperOrigin-RevId: 439439106

Author: MarkDaoust

tools/tensorflow_docs/api_generator/config.py

@@ -17,8 +17,9 @@
 class ParserConfig(object):
   """Stores all indexes required to parse the docs."""
 
-  def __init__(self, reference_resolver, duplicates, duplicate_of, tree, index,
-               reverse_index, path_tree, base_dir, code_url_prefix):
+  def __init__(self, *, reference_resolver, duplicates, duplicate_of, tree,
+               index, reverse_index, path_tree, api_tree, base_dir,
+               code_url_prefix):
     """Object with the common config for docs_for_object() calls.
 
     Args:
@@ -32,6 +33,8 @@ def __init__(self, reference_resolver, duplicates, duplicate_of, tree, index,
         members. Used to populate the members section of a class or module page.
       index: A `dict` mapping full names to objects.
       reverse_index: A `dict` mapping object ids to full names.
+      path_tree: A PathTree datastructure to manage all the API paths.
+      api_tree: A PathTree datastructure to manage all the API objects.
       base_dir: A base path that is stripped from file locations written to the
         docs.
       code_url_prefix: A Url to pre-pend to the links to file locations.
@@ -43,6 +46,7 @@ def __init__(self, reference_resolver, duplicates, duplicate_of, tree, index,
     self.reverse_index = reverse_index
     self.index = index
     self.path_tree = path_tree
+    self.api_tree = api_tree
     self.base_dir = base_dir
     self.code_url_prefix = code_url_prefix
 

tools/tensorflow_docs/api_generator/doc_generator_visitor.py

@@ -16,7 +16,7 @@
 
 import collections
 import dataclasses
-import functools
+import enum
 import inspect
 
 from typing import Any, Dict, List, Optional, NamedTuple, Tuple
@@ -182,6 +182,7 @@ def __init__(self):
     self._duplicate_of: Dict[str, str] = None
 
     self.path_tree = PathTree()
+    self.api_tree = None
 
   @property
   def index(self):
@@ -390,6 +391,8 @@ def build(self):
     if self._reverse_index is not None:
       return
 
+    self.api_tree = ApiTree.from_path_tree(self.path_tree, self._score_name)
+
     # Maps the id of a symbol to its fully qualified name. For symbols that have
     # several aliases, this map contains the first one found.
     # We use id(py_object) to get a hashable value for py_object. Note all
@@ -445,12 +448,32 @@ def build(self):
 
 @dataclasses.dataclass(repr=False)
 class ApiTreeNode(PathTreeNode):
+  """A node in the ApiTree."""
   aliases: List[ApiPath] = dataclasses.field(default_factory=list)
 
   @property
   def obj_type(self) -> obj_type_lib.ObjType:
     return obj_type_lib.ObjType.get(self.py_object)
 
+  class OutputType(enum.Enum):
+    PAGE = 'page'
+    FRAGMENT = 'fragment'
+
+  def output_type(self) -> OutputType:
+    obj_type = obj_type_lib.ObjType.get(self.py_object)
+
+    if obj_type in (obj_type_lib.ObjType.CLASS, obj_type_lib.ObjType.MODULE):
+      return self.OutputType.PAGE
+    elif obj_type in (obj_type_lib.ObjType.CALLABLE,
+                      obj_type_lib.ObjType.TYPE_ALIAS):
+      parent_type = obj_type_lib.ObjType.get(self.parent.py_object)
+      if parent_type is obj_type_lib.ObjType.CLASS:
+        return self.OutputType.FRAGMENT
+      else:
+        return self.OutputType.PAGE
+    else:
+      return self.OutputType.FRAGMENT
+
 
 class ApiTree(Dict[ApiPath, ApiTreeNode]):
   """Public API index.