Simplify test object construction.

Manually building the indices is complicated and error prone. Never do it manually.

PiperOrigin-RevId: 438277327

Author: MarkDaoust

tools/tensorflow_docs/api_generator/doc_generator_visitor.py

@@ -205,7 +205,7 @@ def __init__(self):
     self._duplicates: Dict[str, List[str]] = None
     self._duplicate_of: Dict[str, str] = None
 
-    self._path_tree = PathTree()
+    self.path_tree = PathTree()
 
   @property
   def index(self):
@@ -242,7 +242,6 @@ def reverse_index(self):
     Returns:
       The `id(object)` to full name map.
     """
-    self._maybe_find_duplicates()
     return self._reverse_index
 
   @property
@@ -256,7 +255,6 @@ def duplicate_of(self):
     Returns:
       The map from duplicate name to preferred name.
     """
-    self._maybe_find_duplicates()
     return self._duplicate_of
 
   @property
@@ -273,7 +271,6 @@ def duplicates(self):
     Returns:
       The map from main name to list of all duplicate names.
     """
-    self._maybe_find_duplicates()
     return self._duplicates
 
   def __call__(self, parent_path, parent, children):
@@ -301,16 +298,16 @@ class or module.
     parent_name = '.'.join(parent_path)
     self._index[parent_name] = parent
     self._tree[parent_name] = []
-    if parent_path not in self._path_tree:
-      self._path_tree[parent_path] = parent
+    if parent_path not in self.path_tree:
+      self.path_tree[parent_path] = parent
 
     if not (inspect.ismodule(parent) or inspect.isclass(parent)):
       raise TypeError('Unexpected type in visitor -- '
                       f'{parent_name}: {parent!r}')
 
     for name, child in children:
       child_path = parent_path + (name,)
-      self._path_tree[child_path] = child
+      self.path_tree[child_path] = child
 
       full_name = '.'.join([parent_name, name]) if parent_name else name
       self._index[full_name] = child
@@ -379,7 +376,7 @@ def _score_name(self, name):
     return (defining_class_score, experimental_score, keras_score,
             module_length_score, name)
 
-  def _maybe_find_duplicates(self):
+  def build(self):
     """Compute data structures containing information about duplicates.
 
     Find duplicates in `index` and decide on one to be the "main" name.
@@ -411,7 +408,7 @@ def _maybe_find_duplicates(self):
     # symbol (incl. itself).
     duplicates = {}
 
-    for path, node in self._path_tree.items():
+    for path, node in self.path_tree.items():
       if not path:
         continue
       full_name = node.full_name
@@ -420,7 +417,7 @@ def _maybe_find_duplicates(self):
       if full_name in duplicates:
         continue
 
-      aliases = self._path_tree.nodes_for_obj(py_object)
+      aliases = self.path_tree.nodes_for_obj(py_object)
       # maybe_singleton types can't be looked up by object.
       if not aliases:
         aliases = [node]

tools/tensorflow_docs/api_generator/generate_lib.py

@@ -39,26 +39,6 @@
 
 import yaml
 
-try:
-  # TODO(markdaoust) delete this when the warning is in a stable release.
-  _estimator = importlib.import_module(
-      'tensorflow_estimator.python.estimator.estimator')
-
-  if doc_controls.get_inheritable_header(_estimator.Estimator) is None:
-    _add_header = doc_controls.inheritable_header("""\
-        Warning: Estimators are not recommended for new code.  Estimators run
-        `v1.Session`-style code which is more difficult to write correctly, and
-        can behave unexpectedly, especially when combined with TF 2 code.
-        Estimators do fall under our
-        [compatibility guarantees](https://tensorflow.org/guide/versions), but
-        will receive no fixes other than security vulnerabilities. See the
-        [migration guide](https://tensorflow.org/guide/migrate) for details.
-        """)
-    _add_header(_estimator.Estimator)
-except ImportError:
-  pass
-
-
 # Used to add a collections.OrderedDict representer to yaml so that the
 # dump doesn't contain !!OrderedDict yaml tags.
 # Reference: https://stackoverflow.com/a/21048064
@@ -687,6 +667,7 @@ def extract(py_modules,
 
   traverse.traverse(py_module, visitors, short_name)
 
+  accumulator.build()
   return accumulator
 
 
@@ -818,12 +799,17 @@ def run_extraction(self):
 
     Returns:
     """
-    return extract(
+    visitor = extract(
         py_modules=self._py_modules,
         base_dir=self._base_dir,
         private_map=self._private_map,
         visitor_cls=self._visitor_cls,
         callbacks=self._callbacks)
+    reference_resolver = self.make_reference_resolver(visitor)
+
+    # Write the api docs.
+    parser_config = self.make_parser_config(visitor, reference_resolver)
+    return parser_config
 
   def build(self, output_dir):
     """Build all the docs.
@@ -838,11 +824,7 @@ def build(self, output_dir):
     workdir = pathlib.Path(tempfile.mkdtemp())
 
     # Extract the python api from the _py_modules
-    visitor = self.run_extraction()
-    reference_resolver = self.make_reference_resolver(visitor)
-
-    # Write the api docs.
-    parser_config = self.make_parser_config(visitor, reference_resolver)
+    parser_config = self.run_extraction()
     work_py_dir = workdir / 'api_docs/python'
     write_docs(
         output_dir=str(work_py_dir),
@@ -859,7 +841,7 @@ def build(self, output_dir):
     )
 
     if self.api_cache:
-      reference_resolver.to_json_file(
+      parser_config.reference_resolver.to_json_file(
           str(work_py_dir / self._short_name.replace('.', '/') /
               '_api_cache.json'))
 

tools/tensorflow_docs/api_generator/generate_lib_test.py

@@ -19,6 +19,7 @@
 import sys
 import tempfile
 import textwrap
+import types
 
 from absl import flags
 from absl.testing import absltest
@@ -77,59 +78,21 @@ def setUp(self):
   def get_test_objects(self):
     # These are all mutable objects, so rebuild them for each test.
     # Don't cache the objects.
-    module = sys.modules[__name__]
-
-    index = {
-        'tf':
-            sys,  # Can be any module, this test doesn't care about content.
-        'tf.TestModule':
-            module,
-        'tf.test_function':
-            test_function,
-        'tf.TestModule.test_function':
-            test_function,
-        'tf.TestModule.TestClass':
-            TestClass,
-        'tf.TestModule.TestClass.ChildClass':
-            TestClass.ChildClass,
-        'tf.TestModule.TestClass.ChildClass.GrandChildClass':
-            TestClass.ChildClass.GrandChildClass,
-    }
-
-    tree = {
-        'tf': ['TestModule', 'test_function'],
-        'tf.TestModule': ['test_function', 'TestClass'],
-        'tf.TestModule.TestClass': ['ChildClass'],
-        'tf.TestModule.TestClass.ChildClass': ['GrandChildClass'],
-        'tf.TestModule.TestClass.ChildClass.GrandChildClass': []
-    }
-
-    duplicate_of = {'tf.test_function': 'tf.TestModule.test_function'}
-
-    duplicates = {
-        'tf.TestModule.test_function': [
-            'tf.test_function', 'tf.TestModule.test_function'
-        ]
-    }
-
-    base_dir = os.path.dirname(__file__)
-
-    visitor = DummyVisitor(index, duplicate_of)
-
-    reference_resolver = reference_resolver_lib.ReferenceResolver.from_visitor(
-        visitor=visitor, py_module_names=['tf'], link_prefix='api_docs/python')
-
-    parser_config = config.ParserConfig(
-        reference_resolver=reference_resolver,
-        duplicates=duplicates,
-        duplicate_of=duplicate_of,
-        tree=tree,
-        index=index,
-        reverse_index={},
-        base_dir=base_dir,
-        code_url_prefix='/')
-
-    return reference_resolver, parser_config
+    tf = types.ModuleType('tf')
+    tf.__file__ = __file__
+    tf.TestModule = types.ModuleType('module')
+    tf.test_function = test_function
+    tf.TestModule.test_function = test_function
+    tf.TestModule.TestClass = TestClass
+
+    generator = generate_lib.DocGenerator(
+        root_title='TensorFlow',
+        py_modules=[('tf', tf)],
+        code_url_prefix='https://tensorflow.org/')
+
+    parser_config = generator.run_extraction()
+
+    return parser_config.reference_resolver, parser_config
 
   def test_write(self):
     _, parser_config = self.get_test_objects()