Pure refactor: Give the PageInfo access to the parser_config and api_node.

I prefer api_nodes compared to passing around (py_object, full_name) everywhere.

PiperOrigin-RevId: 444817172

Author: MarkDaoust

tools/tensorflow_docs/api_generator/doc_generator_visitor.py

@@ -73,7 +73,7 @@ class PathTreeNode(object):
   """
   path: ApiPath
   py_object: Any
-  parent: Optional['PathTreeNode']
+  parent: Optional['PathTreeNode'] = None
   children: Dict[str, 'PathTreeNode'] = dataclasses.field(default_factory=dict)
 
   def __hash__(self):

tools/tensorflow_docs/api_generator/generate_lib.py

@@ -120,18 +120,17 @@ def write_docs(
 
   # Parse and write Markdown pages, resolving cross-links (`tf.symbol`).
   num_docs_output = 0
-  for node in parser_config.api_tree.iter_nodes():
-    full_name = node.full_name
-    py_object = node.py_object
+  for api_node in parser_config.api_tree.iter_nodes():
+    full_name = api_node.full_name
+    py_object = api_node.py_object
 
-    if node.output_type() is node.OutputType.FRAGMENT:
+    if api_node.output_type() is api_node.OutputType.FRAGMENT:
       continue
 
     # Generate docs for `py_object`, resolving references.
     try:
       page_info = docs_for_object.docs_for_object(
-          full_name=full_name,
-          py_object=py_object,
+          api_node=api_node,
           parser_config=parser_config,
           extra_docs=extra_docs,
           search_hints=search_hints,

tools/tensorflow_docs/api_generator/parser.py

@@ -33,25 +33,6 @@
 from tensorflow_docs.api_generator import signature as signature_lib
 
 
-@dataclasses.dataclass
-class FileLocation(object):
-  """This class indicates that the object is defined in a regular file.
-
-  This can be used for the `defined_in` slot of the `PageInfo` objects.
-  """
-
-  base_url: Optional[str] = None
-  start_line: Optional[int] = None
-  end_line: Optional[int] = None
-
-  @property
-  def url(self) -> Optional[str]:
-    if self.start_line and self.end_line:
-      if 'github.com' in self.base_url:
-        return f'{self.base_url}#L{self.start_line}-L{self.end_line}'
-    return self.base_url
-
-
 def is_class_attr(full_name, index):
   """Check if the object's parent is a class.
 
@@ -614,6 +595,24 @@ def _unwrap_obj(obj):
   return obj
 
 
+@dataclasses.dataclass
+class FileLocation(object):
+  """This class indicates that the object is defined in a regular file.
+
+  This can be used for the `defined_in` slot of the `PageInfo` objects.
+  """
+
+  base_url: Optional[str] = None
+  start_line: Optional[int] = None
+  end_line: Optional[int] = None
+
+  @property
+  def url(self) -> Optional[str]:
+    if self.start_line and self.end_line:
+      if 'github.com' in self.base_url:
+        return f'{self.base_url}#L{self.start_line}-L{self.end_line}'
+    return self.base_url
+
 def get_defined_in(
     py_object: Any,
     parser_config: config.ParserConfig) -> Optional[FileLocation]:

tools/tensorflow_docs/api_generator/parser_test.py

@@ -28,6 +28,7 @@
 
 from tensorflow_docs.api_generator import config
 from tensorflow_docs.api_generator import doc_controls
+from tensorflow_docs.api_generator import doc_generator_visitor
 from tensorflow_docs.api_generator import generate_lib
 from tensorflow_docs.api_generator import parser
 from tensorflow_docs.api_generator import reference_resolver as reference_resolver_lib
@@ -198,10 +199,13 @@ def test_docs_for_class(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=(
+            'm',
+            'TestClass',
+        ), py_object=TestClass)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.TestClass',
-        py_object=TestClass,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -241,10 +245,10 @@ def test_dataclass_attributes_table(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m', 'ExampleDataclass'), py_object=ExampleDataclass)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.ExampleDataclass',
-        py_object=ExampleDataclass,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     self.assertCountEqual(['a', 'b', 'c', 'x', 'y', 'z'],
                           [name for name, value in page_info.attr_block.items])
@@ -268,10 +272,10 @@ def hide(path, parent, children):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m', 'namedtupleclass'), py_object=namedtupleclass)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.namedtupleclass',
-        py_object=namedtupleclass,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     self.assertIsNone(page_info._namedtuplefields['hidden'])
 
@@ -308,8 +312,13 @@ def a_method(self, arg='default'):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=(
+            'm',
+            'Child',
+        ), py_object=Child)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.Child', py_object=Child, parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     # Make sure the `a_method` is not present
     self.assertEmpty(page_info.methods)
@@ -347,10 +356,10 @@ def my_method(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m', 'ChildMessage'), py_object=ChildMessage)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.ChildMessage',
-        py_object=ChildMessage,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     self.assertLen(page_info.methods, 1)
     self.assertEqual('my_method', page_info.methods[0].short_name)
@@ -369,8 +378,10 @@ def test_docs_for_module(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m',), py_object=test_module)
     page_info = docs_for_object.docs_for_object(
-        full_name='m', py_object=test_module, parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -395,10 +406,10 @@ def test_docs_for_function(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('test_function',), py_object=test_function)
     page_info = docs_for_object.docs_for_object(
-        full_name='test_function',
-        py_object=test_function,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -420,10 +431,11 @@ def test_docs_for_function_with_kwargs(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('test_function_with_args_kwargs',),
+        py_object=test_function_with_args_kwargs)
     page_info = docs_for_object.docs_for_object(
-        full_name='test_function_with_args_kwargs',
-        py_object=test_function_with_args_kwargs,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     # Make sure the brief docstring is present
     self.assertEqual(
@@ -594,10 +606,10 @@ def test_getsource_indexerror_resilience(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m', 'ConcreteMutableMapping'), py_object=ConcreteMutableMapping)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.ConcreteMutableMapping',
-        py_object=ConcreteMutableMapping,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     self.assertIn(ConcreteMutableMapping.get,
                   [m.py_object for m in page_info.methods])
@@ -621,8 +633,10 @@ def test_strips_default_arg_memory_address(self):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=('m', 'fun'), py_object=m.fun)
     page_info = docs_for_object.docs_for_object(
-        full_name='m.fun', py_object=m.fun, parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     output = str(page_info.signature)
     self.assertNotIn('object at 0x', output)
@@ -669,10 +683,10 @@ def test_empty_defined_in(self, cls, method, py_object):
 
     parser_config = generator.run_extraction()
 
+    api_node = doc_generator_visitor.ApiTreeNode(
+        path=(cls, method), py_object=py_object)
     function_info = docs_for_object.docs_for_object(
-        full_name='%s.%s' % (cls, method),
-        py_object=py_object,
-        parser_config=parser_config)
+        api_node=api_node, parser_config=parser_config)
 
     self.assertIsNone(function_info.defined_in)
 

tools/tensorflow_docs/api_generator/pretty_docs/base_page.py

@@ -110,10 +110,10 @@ class PageInfo:
 
   def __init__(
       self,
-      full_name: str,
-      py_object: Any,
+      api_node,
       extra_docs: Optional[Dict[int, str]] = None,
       search_hints: bool = True,
+      parser_config=None,
   ):
     """Initialize a PageInfo.
 
@@ -126,46 +126,48 @@ def __init__(
         "robots: noindex"
 
     """
-    self.full_name = full_name
-    self.py_object = py_object
+    self.api_node = api_node
+    self.full_name = api_node.full_name
+    self.py_object = api_node.py_object
     self._extra_docs = extra_docs
     self.search_hints = search_hints
+    self.parser_config = parser_config
 
     self._defined_in = None
     self._aliases = None
     self._doc = None
     self._page_text = None
 
-  def collect_docs(self, parser_config: config.ParserConfig):
+  def collect_docs(self):
     """Collects additional information from the `config.ParserConfig`."""
     pass
 
-  def docs_for_object(self, parser_config):
-    duplicate_names = parser_config.duplicates.get(self.full_name, [])
-    if self.full_name in duplicate_names:
-      duplicate_names.remove(self.full_name)
-
+  def docs_for_object(self):
     relative_path = os.path.relpath(
         path='.',
         start=os.path.dirname(parser.documentation_path(self.full_name)) or '.')
 
     # Convert from OS-specific path to URL/POSIX path.
     relative_path = posixpath.join(*relative_path.split(os.path.sep))
 
-    with parser_config.reference_resolver.temp_prefix(relative_path):
+    with self.parser_config.reference_resolver.temp_prefix(relative_path):
       self.set_doc(
           parser.parse_md_docstring(
               self.py_object,
               self.full_name,
-              parser_config,
+              self.parser_config,
               self._extra_docs,
           ))
 
-      self.collect_docs(parser_config)
+      self.collect_docs()
 
-      self.set_aliases(duplicate_names)
+      aliases = ['.'.join(alias) for alias in self.api_node.aliases]
+      if self.full_name in aliases:
+        aliases.remove(self.full_name)
+      self.set_aliases(aliases)
 
-      self.set_defined_in(parser.get_defined_in(self.py_object, parser_config))
+      self.set_defined_in(
+          parser.get_defined_in(self.py_object, self.parser_config))
 
       self._page_text = self.build()