Fix TensorFlow docs cannonical name chooser

Prefer `tf.submodule.thing` to `tf.thing` as the canonical doc location.

For class members, prefer the defining class.

Always prefer the non-contrib version.

Choose the shortest name, but prefer sub-modules.
(1.10 is includes submodules like tf.io and tf.math)

PiperOrigin-RevId: 207756817

4 files changed, 249 insertions, 50 deletions

diff --git a/tensorflow/tools/docs/BUILD b/tensorflow/tools/docs/BUILD
index 66b10478ac..cc7885ab1b 100644
--- a/tensorflow/tools/docs/BUILD
+++ b/tensorflow/tools/docs/BUILD
@@ -28,6 +28,7 @@ py_test(
     srcs_version = "PY2AND3",
     deps = [
         ":doc_generator_visitor",
+        ":generate_lib",
         "//tensorflow/python:platform_test",
     ],
 )
diff --git a/tensorflow/tools/docs/doc_generator_visitor.py b/tensorflow/tools/docs/doc_generator_visitor.py
index c090dbd8da..e5eaf8cc05 100644
--- a/tensorflow/tools/docs/doc_generator_visitor.py
+++ b/tensorflow/tools/docs/doc_generator_visitor.py
@@ -159,6 +159,55 @@ class DocGeneratorVisitor(object):
       self._index[full_name] = child
       self._tree[parent_name].append(name)
 
+  def _score_name(self, name):
+    """Return a tuple of scores indicating how to sort for the best name.
+
+    This function is meant to be used as the `key` to the `sorted` function.
+
+    This sorting in order:
+      Prefers names refering to the defining class, over a subclass.
+      Prefers names that are not in "contrib".
+      prefers submodules to the root namespace.
+      Prefers short names `tf.thing` over `tf.a.b.c.thing`
+      Sorts lexicographically on name parts.
+
+    Args:
+      name: the full name to score, for example `tf.estimator.Estimator`
+
+    Returns:
+      A tuple of scores. When sorted the preferred name will have the lowest
+      value.
+    """
+    parts = name.split('.')
+    short_name = parts[-1]
+
+    container = self._index['.'.join(parts[:-1])]
+
+    defining_class_score = 1
+    if tf_inspect.isclass(container):
+      if short_name in container.__dict__:
+        # prefer the defining class
+        defining_class_score = -1
+
+    contrib_score = -1
+    if 'contrib' in parts:
+      contrib_score = 1
+
+    while parts:
+      parts.pop()
+      container = self._index['.'.join(parts)]
+      if tf_inspect.ismodule(container):
+        break
+    module_length = len(parts)
+    if len(parts) == 2:
+      # `tf.submodule.thing` is better than `tf.thing`
+      module_length_score = -1
+    else:
+      # shorter is better
+      module_length_score = module_length
+
+    return (defining_class_score, contrib_score, module_length_score, name)
+
   def _maybe_find_duplicates(self):
     """Compute data structures containing information about duplicates.
 
@@ -192,7 +241,7 @@ class DocGeneratorVisitor(object):
       if (py_object is not None and
           not isinstance(py_object, six.integer_types + six.string_types +
                          (six.binary_type, six.text_type, float, complex, bool))
-          and py_object is not ()):
+          and py_object is not ()):  # pylint: disable=literal-comparison
         object_id = id(py_object)
         if object_id in reverse_index:
           master_name = reverse_index[object_id]
@@ -217,9 +266,10 @@ class DocGeneratorVisitor(object):
       if master_name:
         master_name = 'tf.%s' % master_name
       else:
-        # Choose the lexicographically first name with the minimum number of
-        # submodules. This will prefer highest level namespace for any symbol.
-        master_name = min(names, key=lambda name: name.count('.'))
+        # Choose the master name with a lexical sort on the tuples returned by
+        # by _score_name.
+        master_name = min(names, key=self._score_name)
+        print(names, master_name)
 
       duplicates[master_name] = names
       for name in names:
diff --git a/tensorflow/tools/docs/doc_generator_visitor_test.py b/tensorflow/tools/docs/doc_generator_visitor_test.py
index cf5be45f40..1c2635d4a8 100644
--- a/tensorflow/tools/docs/doc_generator_visitor_test.py
+++ b/tensorflow/tools/docs/doc_generator_visitor_test.py
@@ -18,8 +18,21 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import types
+
 from tensorflow.python.platform import googletest
 from tensorflow.tools.docs import doc_generator_visitor
+from tensorflow.tools.docs import generate_lib
+
+
+class NoDunderVisitor(doc_generator_visitor.DocGeneratorVisitor):
+
+  def __call__(self, parent_name, parent, children):
+    """Drop all the dunder methods to make testing easier."""
+    children = [
+        (name, obj) for (name, obj) in children if not name.startswith('_')
+    ]
+    super(NoDunderVisitor, self).__call__(parent_name, parent, children)
 
 
 class DocGeneratorVisitorTest(googletest.TestCase):
@@ -57,52 +70,184 @@ class DocGeneratorVisitorTest(googletest.TestCase):
     with self.assertRaises(RuntimeError):
       visitor('non_class_or_module', 'non_class_or_module_object', [])
 
-  def test_duplicates(self):
-    visitor = doc_generator_visitor.DocGeneratorVisitor()
-    visitor(
-        'submodule.DocGeneratorVisitor',
-        doc_generator_visitor.DocGeneratorVisitor,
-        [('index', doc_generator_visitor.DocGeneratorVisitor.index),
-         ('index2', doc_generator_visitor.DocGeneratorVisitor.index)])
-    visitor(
-        'submodule2.DocGeneratorVisitor',
-        doc_generator_visitor.DocGeneratorVisitor,
-        [('index', doc_generator_visitor.DocGeneratorVisitor.index),
-         ('index2', doc_generator_visitor.DocGeneratorVisitor.index)])
-    visitor(
-        'DocGeneratorVisitor2',
-        doc_generator_visitor.DocGeneratorVisitor,
-        [('index', doc_generator_visitor.DocGeneratorVisitor.index),
-         ('index2', doc_generator_visitor.DocGeneratorVisitor.index)])
-
-    # The shorter path should be master, or if equal, the lexicographically
-    # first will be.
-    self.assertEqual(
-        {'DocGeneratorVisitor2': sorted(['submodule.DocGeneratorVisitor',
-                                         'submodule2.DocGeneratorVisitor',
-                                         'DocGeneratorVisitor2']),
-         'DocGeneratorVisitor2.index': sorted([
-             'submodule.DocGeneratorVisitor.index',
-             'submodule.DocGeneratorVisitor.index2',
-             'submodule2.DocGeneratorVisitor.index',
-             'submodule2.DocGeneratorVisitor.index2',
-             'DocGeneratorVisitor2.index',
-             'DocGeneratorVisitor2.index2'
-         ]),
-        }, visitor.duplicates)
-    self.assertEqual({
-        'submodule.DocGeneratorVisitor': 'DocGeneratorVisitor2',
-        'submodule.DocGeneratorVisitor.index': 'DocGeneratorVisitor2.index',
-        'submodule.DocGeneratorVisitor.index2': 'DocGeneratorVisitor2.index',
-        'submodule2.DocGeneratorVisitor': 'DocGeneratorVisitor2',
-        'submodule2.DocGeneratorVisitor.index': 'DocGeneratorVisitor2.index',
-        'submodule2.DocGeneratorVisitor.index2': 'DocGeneratorVisitor2.index',
-        'DocGeneratorVisitor2.index2': 'DocGeneratorVisitor2.index'
+  def test_duplicates_module_class_depth(self):
+
+    class Parent(object):
+
+      class Nested(object):
+        pass
+
+    tf = types.ModuleType('tf')
+    tf.Parent = Parent
+    tf.submodule = types.ModuleType('submodule')
+    tf.submodule.Parent = Parent
+
+    visitor = generate_lib.extract(
+        [('tf', tf)],
+        private_map={},
+        do_not_descend_map={},
+        visitor_cls=NoDunderVisitor)
+
+    self.assertEqual({
+        'tf.submodule.Parent':
+            sorted([
+                'tf.Parent',
+                'tf.submodule.Parent',
+            ]),
+        'tf.submodule.Parent.Nested':
+            sorted([
+                'tf.Parent.Nested',
+                'tf.submodule.Parent.Nested',
+            ]),
+    }, visitor.duplicates)
+
+    self.assertEqual({
+        'tf.Parent.Nested': 'tf.submodule.Parent.Nested',
+        'tf.Parent': 'tf.submodule.Parent',
+    }, visitor.duplicate_of)
+
+    self.assertEqual({
+        id(Parent): 'tf.submodule.Parent',
+        id(Parent.Nested): 'tf.submodule.Parent.Nested',
+        id(tf): 'tf',
+        id(tf.submodule): 'tf.submodule',
+    }, visitor.reverse_index)
+
+  def test_duplicates_contrib(self):
+
+    class Parent(object):
+      pass
+
+    tf = types.ModuleType('tf')
+    tf.contrib = types.ModuleType('contrib')
+    tf.submodule = types.ModuleType('submodule')
+    tf.contrib.Parent = Parent
+    tf.submodule.Parent = Parent
+
+    visitor = generate_lib.extract(
+        [('tf', tf)],
+        private_map={},
+        do_not_descend_map={},
+        visitor_cls=NoDunderVisitor)
+
+    self.assertEqual({
+        'tf.submodule.Parent':
+            sorted(['tf.contrib.Parent', 'tf.submodule.Parent']),
+    }, visitor.duplicates)
+
+    self.assertEqual({
+        'tf.contrib.Parent': 'tf.submodule.Parent',
+    }, visitor.duplicate_of)
+
+    self.assertEqual({
+        id(tf): 'tf',
+        id(tf.submodule): 'tf.submodule',
+        id(Parent): 'tf.submodule.Parent',
+        id(tf.contrib): 'tf.contrib',
+    }, visitor.reverse_index)
+
+  def test_duplicates_defining_class(self):
+
+    class Parent(object):
+      obj1 = object()
+
+    class Child(Parent):
+      pass
+
+    tf = types.ModuleType('tf')
+    tf.Parent = Parent
+    tf.Child = Child
+
+    visitor = generate_lib.extract(
+        [('tf', tf)],
+        private_map={},
+        do_not_descend_map={},
+        visitor_cls=NoDunderVisitor)
+
+    self.assertEqual({
+        'tf.Parent.obj1': sorted([
+            'tf.Parent.obj1',
+            'tf.Child.obj1',
+        ]),
+    }, visitor.duplicates)
+
+    self.assertEqual({
+        'tf.Child.obj1': 'tf.Parent.obj1',
     }, visitor.duplicate_of)
+
+    self.assertEqual({
+        id(tf): 'tf',
+        id(Parent): 'tf.Parent',
+        id(Child): 'tf.Child',
+        id(Parent.obj1): 'tf.Parent.obj1',
+    }, visitor.reverse_index)
+
+  def test_duplicates_module_depth(self):
+
+    class Parent(object):
+      pass
+
+    tf = types.ModuleType('tf')
+    tf.submodule = types.ModuleType('submodule')
+    tf.submodule.submodule2 = types.ModuleType('submodule2')
+    tf.Parent = Parent
+    tf.submodule.submodule2.Parent = Parent
+
+    visitor = generate_lib.extract(
+        [('tf', tf)],
+        private_map={},
+        do_not_descend_map={},
+        visitor_cls=NoDunderVisitor)
+
+    self.assertEqual({
+        'tf.Parent': sorted(['tf.Parent', 'tf.submodule.submodule2.Parent']),
+    }, visitor.duplicates)
+
+    self.assertEqual({
+        'tf.submodule.submodule2.Parent': 'tf.Parent'
+    }, visitor.duplicate_of)
+
+    self.assertEqual({
+        id(tf): 'tf',
+        id(tf.submodule): 'tf.submodule',
+        id(tf.submodule.submodule2): 'tf.submodule.submodule2',
+        id(Parent): 'tf.Parent',
+    }, visitor.reverse_index)
+
+  def test_duplicates_name(self):
+
+    class Parent(object):
+      obj1 = object()
+
+    Parent.obj2 = Parent.obj1
+
+    tf = types.ModuleType('tf')
+    tf.submodule = types.ModuleType('submodule')
+    tf.submodule.Parent = Parent
+
+    visitor = generate_lib.extract(
+        [('tf', tf)],
+        private_map={},
+        do_not_descend_map={},
+        visitor_cls=NoDunderVisitor)
+
+    self.assertEqual({
+        'tf.submodule.Parent.obj1':
+            sorted([
+                'tf.submodule.Parent.obj1',
+                'tf.submodule.Parent.obj2',
+            ]),
+    }, visitor.duplicates)
+
+    self.assertEqual({
+        'tf.submodule.Parent.obj2': 'tf.submodule.Parent.obj1',
+    }, visitor.duplicate_of)
+
     self.assertEqual({
-        id(doc_generator_visitor.DocGeneratorVisitor): 'DocGeneratorVisitor2',
-        id(doc_generator_visitor.DocGeneratorVisitor.index):
-        'DocGeneratorVisitor2.index',
+        id(tf): 'tf',
+        id(tf.submodule): 'tf.submodule',
+        id(Parent): 'tf.submodule.Parent',
+        id(Parent.obj1): 'tf.submodule.Parent.obj1',
     }, visitor.reverse_index)
 
 if __name__ == '__main__':
diff --git a/tensorflow/tools/docs/generate_lib.py b/tensorflow/tools/docs/generate_lib.py
index 4f70a69364..7c8dbd5dc2 100644
--- a/tensorflow/tools/docs/generate_lib.py
+++ b/tensorflow/tools/docs/generate_lib.py
@@ -284,10 +284,13 @@ def _get_default_do_not_descend_map():
   }
 
 
-def extract(py_modules, private_map, do_not_descend_map):
+def extract(py_modules,
+            private_map,
+            do_not_descend_map,
+            visitor_cls=doc_generator_visitor.DocGeneratorVisitor):
   """Extract docs from tf namespace and write them to disk."""
   # Traverse the first module.
-  visitor = doc_generator_visitor.DocGeneratorVisitor(py_modules[0][0])
+  visitor = visitor_cls(py_modules[0][0])
   api_visitor = public_api.PublicAPIVisitor(visitor)
   api_visitor.set_root_name(py_modules[0][0])
   add_dict_to_dict(private_map, api_visitor.private_map)