Remove TF1 only stuff from generate2.py

PiperOrigin-RevId: 286656789
Change-Id: I1959fd8cbf368efc7c9647d0ca101704100ebd62

tensorflow/tools/docs/generate2.py

@@ -20,9 +20,9 @@ python generate2.py --output_dir=/tmp/out
 
 Requires a local installation of `tensorflow_docs`:
 
-  ```
+```
-  pip install git+https://github.com/tensorflow/docs
+pip install git+https://github.com/tensorflow/docs
-  ```
+```
 """
 
 from __future__ import absolute_import
@@ -34,7 +34,6 @@ import textwrap
 
 from absl import app
 from absl import flags
-from distutils.version import LooseVersion
 
 import tensorflow as tf
 
@@ -56,7 +55,6 @@ parser.tf_inspect = tf_inspect
 # So patch `tf.__all__` to list everything.
 tf.__all__ = [item_name for item_name, value in tf_inspect.getmembers(tf)]
 
-
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string(
@@ -64,41 +62,25 @@ flags.DEFINE_string(
     "/code/stable/tensorflow",
     "A url to prepend to code paths when creating links to defining code")
 
-flags.DEFINE_string(
+flags.DEFINE_string("output_dir", "/tmp/out",
-    "output_dir", "/tmp/out",
                     "A directory, where the docs will be output to.")
 
 flags.DEFINE_bool("search_hints", True,
                   "Include meta-data search hints at the top of each file.")
 
-flags.DEFINE_string("site_path", "",
+flags.DEFINE_string(
-                    "The prefix ({site-path}/api_docs/python/...) used in the "
+    "site_path", "", "The prefix ({site-path}/api_docs/python/...) used in the "
     "`_toc.yaml` and `_redirects.yaml` files")
 
-
+_PRIVATE_MAP = {
-if tf.__version__.startswith('1'):
+    "tf": ["python", "core", "compiler", "examples", "tools"],
-  PRIVATE_MAP = {
-      'tf.test': ['mock'],
-      'tf': ['python', 'core', 'compiler', 'examples', 'tools', 'contrib'],
     # There's some aliasing between the compats and v1/2s, so it's easier to
     # block by name and location than by deleting, or hiding objects.
-      'tf.compat.v1.compat': ['v1', 'v2'],
+    "tf.compat.v1.compat": ["v1", "v2"],
-      'tf.compat.v2.compat': ['v1', 'v2']
+    "tf.compat.v2.compat": ["v1", "v2"]
-  }
+}
 
-  DO_NOT_DESCEND_MAP = {
+tf.__doc__ = """
-      'tf': ['cli', 'lib', 'wrappers', 'contrib'],
-  }
-else:
-  PRIVATE_MAP = {
-      'tf': ['python', 'core', 'compiler', 'examples', 'tools'],
-      # There's some aliasing between the compats and v1/2s, so it's easier to
-      # block by name and location than by deleting, or hiding objects.
-      'tf.compat.v1.compat': ['v1', 'v2'],
-      'tf.compat.v2.compat': ['v1', 'v2']
-  }
-  DO_NOT_DESCEND_MAP = {}
-  tf.__doc__ = """
   ## TensorFlow
 
   ```
@@ -112,27 +94,14 @@ _raw_ops_doc = textwrap.dedent("""\n
   for details. Unless you are library writer, you likely do not need to use these
   ops directly.""")
 
-if LooseVersion(tf.__version__) < LooseVersion('2'):
+tf.raw_ops.__doc__ += _raw_ops_doc
-  tf.raw_ops.__doc__ = _raw_ops_doc
-  tf.contrib.__doc__ = """
-    Contrib module containing volatile or experimental code.
-
-    Warning: The `tf.contrib` module will not be included in TensorFlow 2.0. Many
-    of its submodules have been integrated into TensorFlow core, or spun-off into
-    other projects like [`tensorflow_io`](https://github.com/tensorflow/io), or
-    [`tensorflow_addons`](https://github.com/tensorflow/addons). For instructions
-    on how to upgrade see the
-    [Migration guide](https://www.tensorflow.org/guide/migrate).
-    """
-else:
-  tf.raw_ops.__doc__ += _raw_ops_doc
 
 
 # The doc generator isn't aware of tf_export.
 # So prefix the score tuples with -1 when this is the canonical name, +1
 # otherwise. The generator chooses the name with the lowest score.
-class TfExportAwareDocGeneratorVisitor(
+class TfExportAwareDocGeneratorVisitor(doc_generator_visitor.DocGeneratorVisitor
-    doc_generator_visitor.DocGeneratorVisitor):
+                                      ):
   """A `tf_export` aware doc_visitor."""
 
   def _score_name(self, name):
@@ -214,28 +183,23 @@ def build_docs(output_dir, code_url_prefix, search_hints=True):
       "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
   )
 
-  if LooseVersion(tf.__version__) < LooseVersion('2'):
-    root_title = 'TensorFlow'
-  elif LooseVersion(tf.__version__) >= LooseVersion('2'):
-    root_title = 'TensorFlow 2.0'
-
   doc_generator = generate_lib.DocGenerator(
-      root_title=root_title,
+      root_title="TensorFlow 2.0",
       py_modules=[("tf", tf)],
       base_dir=base_dirs,
       search_hints=search_hints,
       code_url_prefix=code_url_prefixes,
       site_path=FLAGS.site_path,
       visitor_cls=TfExportAwareDocGeneratorVisitor,
-      private_map=PRIVATE_MAP,
+      private_map=_PRIVATE_MAP)
-      do_not_descend_map=DO_NOT_DESCEND_MAP)
 
   doc_generator.build(output_dir)
 
 
 def main(argv):
   del argv
-  build_docs(output_dir=FLAGS.output_dir,
+  build_docs(
+      output_dir=FLAGS.output_dir,
       code_url_prefix=FLAGS.code_url_prefix,
       search_hints=FLAGS.search_hints)
 

tensorflow/tools/docs/generate2_test.py

@@ -32,6 +32,7 @@ from tensorflow.tools.docs import generate2
 del tf.compat.v2
 del tf.compat.v1
 
+
 class Generate2Test(googletest.TestCase):
 
   def test_end_to_end(self):