Protect against inconsistent path selection.

PiperOrigin-RevId: 442778183

Author: MarkDaoust

tools/tensorflow_docs/api_generator/doc_generator_visitor.py

@@ -19,7 +19,7 @@
 import enum
 import inspect
 
-from typing import Any, Dict, List, Optional, NamedTuple, Tuple
+from typing import Any, Dict, List, Optional, NamedTuple, Sequence, Tuple
 
 from tensorflow_docs.api_generator import obj_type as obj_type_lib
 
@@ -584,7 +584,7 @@ def from_path_tree(cls, path_tree: PathTree, score_name_fn) -> 'ApiTree':
 
       parents = [node.parent for node in duplicate_nodes]
 
-      # Choose the master name with a lexical sort on the tuples returned by
+      # Choose the priority name with a lexical sort on the tuples returned by
       # by _score_name.
       if not all(parent.path in self for parent in parents):
         # rewind
@@ -596,13 +596,34 @@ def from_path_tree(cls, path_tree: PathTree, score_name_fn) -> 'ApiTree':
           active_nodes.appendleft(parent)
         continue
       # If we've made it here, the immediate parents of each of the paths have
-      # been processed, so now we can choose its master name.
+      # been processed, so now we can choose its priority name.
       aliases = [node.path for node in duplicate_nodes]
 
-      master_path = min(aliases, key=score_name_fn)
+      priority_path = self._choose_priority_path(aliases, score_name_fn)
 
-      self.insert(master_path, current_node.py_object, aliases)
+      if priority_path is None:
+        # How did this happen?
+        # No parents in the public api -> you are not in the public API.
+        continue
+
+      self.insert(priority_path, current_node.py_object, aliases)
 
       active_nodes.extend(current_node.children.values())
 
     return self
+
+  def _choose_priority_path(self, aliases: Sequence[ApiPath],
+                            score_name_fn) -> Optional[ApiPath]:
+    # Only consider a path an option for the priority_path if its parent-path
+    # is the priority_path for that object.
+    priority_path_options = []
+    for alias in aliases:
+      parent_path = alias[:-1]
+
+      if self[parent_path].path == parent_path:
+        priority_path_options.append(alias)
+
+    try:
+      return min(priority_path_options, key=score_name_fn)
+    except ValueError:
+      return None

tools/tensorflow_docs/api_generator/doc_generator_visitor_test.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 """Tests for tools.docs.doc_generator_visitor."""
 
-import argparse
 import inspect
 import io
 import os
@@ -262,8 +261,8 @@ class Parent(object):
 class PathTreeTest(absltest.TestCase):
 
   def test_contains(self):
-    tf = argparse.Namespace()
-    tf.sub = argparse.Namespace()
+    tf = types.ModuleType('tf')
+    tf.sub = types.ModuleType('sub')
 
     tree = doc_generator_visitor.PathTree()
     tree[('tf',)] = tf
@@ -273,8 +272,8 @@ def test_contains(self):
     self.assertIn(('tf', 'sub'), tree)
 
   def test_node_insertion(self):
-    tf = argparse.Namespace()
-    tf.sub = argparse.Namespace()
+    tf = types.ModuleType('tf')
+    tf.sub = types.ModuleType('sub')
     tf.sub.object = object()
 
     tree = doc_generator_visitor.PathTree()
@@ -290,10 +289,10 @@ def test_node_insertion(self):
     self.assertIs(node.children['thing'], tree[('tf', 'sub', 'thing')])
 
   def test_duplicate(self):
-    tf = argparse.Namespace()
-    tf.sub = argparse.Namespace()
+    tf = types.ModuleType('tf')
+    tf.sub = types.ModuleType('sub')
     tf.sub.thing = object()
-    tf.sub2 = argparse.Namespace()
+    tf.sub2 = types.ModuleType('sub2')
     tf.sub2.thing = tf.sub.thing
 
     tree = doc_generator_visitor.PathTree()
@@ -308,10 +307,10 @@ def test_duplicate(self):
         [tree[('tf', 'sub', 'thing')], tree[('tf', 'sub2', 'thing')]])
 
   def test_duplicate_singleton(self):
-    tf = argparse.Namespace()
-    tf.sub = argparse.Namespace()
+    tf = types.ModuleType('tf')
+    tf.sub = types.ModuleType('sub')
     tf.sub.thing = 999
-    tf.sub2 = argparse.Namespace()
+    tf.sub2 = types.ModuleType('sub2')
     tf.sub2.thing = tf.sub.thing
 
     tree = doc_generator_visitor.PathTree()
@@ -322,8 +321,7 @@ def test_duplicate_singleton(self):
     tree[('tf', 'sub2', 'thing')] = tf.sub2.thing
 
     found = tree.nodes_for_obj(tf.sub.thing)
-    self.assertIsNotNone(found)
-    self.assertEmpty(found)
+    self.assertEqual([], found)
 
 
 class ApiTreeTest(absltest.TestCase):
@@ -469,6 +467,42 @@ def test_api_tree_toc_integration(self):
 
     self.assertEqual(expected, stream.getvalue())
 
+  def test_non_priority_name(self):
+
+    class Class1:
+      pass
+
+    mod = types.ModuleType('mod')
+    mod.a = types.ModuleType('sub')
+    mod.a.Class1 = Class1
+    mod.b = mod.a
+
+    path_tree = doc_generator_visitor.PathTree()
+    path_tree[('mod',)] = mod
+    path_tree[('mod', 'a')] = mod.a
+    path_tree[('mod', 'a', 'Class1')] = mod.a.Class1
+    path_tree[('mod', 'b')] = mod.b
+    path_tree[('mod', 'b', 'Class1')] = mod.b.Class1
+
+    def inconsistent_name_score(path):
+      # `mod.a` is prefered over `mod.b`, but `b.Class1` is prefered over
+      # `a.Class1`!
+      scores = {
+          ('mod',): 0,
+          ('mod', 'a'): 0,  # prefer 'a'
+          ('mod', 'b'): 1,
+          ('mod', 'a', 'Class1'): 1,
+          ('mod', 'b', 'Class1'): 0,  # prefer 'b.Class1'
+      }
+      return scores[path]
+
+    api_tree = doc_generator_visitor.ApiTree.from_path_tree(
+        path_tree, inconsistent_name_score)
+    node = api_tree.node_for_object(Class1)
+
+    # `Class1` can't choose `b.Class1` as its priority_path because
+    # `a` is the priority_path for `sub`.
+    self.assertEqual('mod.a.Class1', node.full_name)
 
 if __name__ == '__main__':
   absltest.main()