Remove TF1 only stuff from generate2.py

PiperOrigin-RevId: 286656789 Change-Id: I1959fd8cbf368efc7c9647d0ca101704100ebd62
frreiss · Dec 21, 2019 · 8ac21df · 8ac21df
1 parent 0bdc8d1
commit 8ac21df
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 69 deletions.
diff --git a/tensorflow/tools/docs/generate2.py b/tensorflow/tools/docs/generate2.py
@@ -20,9 +20,9 @@
 
 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 @@
 
 from absl import app
 from absl import flags
-from distutils.version import LooseVersion
 
 import tensorflow as tf
 
@@ -56,83 +55,53 @@
 # 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(
     "code_url_prefix",
     "/code/stable/tensorflow",
     "A url to prepend to code paths when creating links to defining code")
 
-flags.DEFINE_string(
-    "output_dir", "/tmp/out",
-    "A directory, where the docs will be output to.")
+flags.DEFINE_string("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", "",
-                    "The prefix ({site-path}/api_docs/python/...) used in the "
-                    "`_toc.yaml` and `_redirects.yaml` files")
-
-
-if tf.__version__.startswith('1'):
-  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.v2.compat': ['v1', 'v2']
-  }
-
-  DO_NOT_DESCEND_MAP = {
-      '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
-
-    ```
-    pip install tensorflow
-    ```
-    """
+flags.DEFINE_string(
+    "site_path", "", "The prefix ({site-path}/api_docs/python/...) used in the "
+    "`_toc.yaml` and `_redirects.yaml` files")
+
+_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"]
+}
+
+tf.__doc__ = """
+  ## TensorFlow
+
+  ```
+  pip install tensorflow
+  ```
+  """
 
 _raw_ops_doc = textwrap.dedent("""\n
   Note: `tf.raw_ops` provides direct/low level access to all TensorFlow ops. See \
   [the RFC](https://github.com/tensorflow/community/blob/master/rfcs/20181225-tf-raw-ops.md)
   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.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
+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(
-    doc_generator_visitor.DocGeneratorVisitor):
+class TfExportAwareDocGeneratorVisitor(doc_generator_visitor.DocGeneratorVisitor
+                                      ):
   """A `tf_export` aware doc_visitor."""
 
   def _score_name(self, name):
@@ -214,30 +183,25 @@ 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,
-      do_not_descend_map=DO_NOT_DESCEND_MAP)
+      private_map=_PRIVATE_MAP)
 
   doc_generator.build(output_dir)
 
 
 def main(argv):
   del argv
-  build_docs(output_dir=FLAGS.output_dir,
-             code_url_prefix=FLAGS.code_url_prefix,
-             search_hints=FLAGS.search_hints)
+  build_docs(
+      output_dir=FLAGS.output_dir,
+      code_url_prefix=FLAGS.code_url_prefix,
+      search_hints=FLAGS.search_hints)
 
 
 if __name__ == "__main__":

diff --git a/tensorflow/tools/docs/generate2_test.py b/tensorflow/tools/docs/generate2_test.py
@@ -32,6 +32,7 @@
 del tf.compat.v2
 del tf.compat.v1
 
+
 class Generate2Test(googletest.TestCase):
 
   def test_end_to_end(self):