Fix doc generator to handle new package layout.

tensorflow_core is gone in tf-nightly.

PiperOrigin-RevId: 296238681
Change-Id: I604be239c807b6e6fb9569560d9f94326b303711

tensorflow/tools/docs/BUILD

@@ -165,11 +165,20 @@ py_binary(
     ],
 )
 
+py_library(
+    # Opensource only
+    name = "base_dir_oss",
+    srcs = ["base_dir.py"],
+    srcs_version = "PY3",
+    deps = [],
+)
+
 py_library(
     name = "generate2_lib",
     srcs = ["generate2.py"],
     srcs_version = "PY3",
     deps = [
+        ":base_dir_oss",
         "//tensorflow:tensorflow_py",
         "//tensorflow/python:util",
     ],

tensorflow/tools/docs/base_dir.py

@@ -0,0 +1,52 @@
+# Lint as: python3
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Opensource base_dir configuration for tensorflow doc-generator."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import distutils
+from os import path
+
+import tensorboard
+import tensorflow as tf
+import tensorflow_estimator
+
+
+def get_base_dirs_and_prefixes(code_url_prefix):
+  """Returns the base_dirs and code_prefixes for OSS TensorFlow api gen."""
+  base_dir = path.dirname(tf.__file__)
+
+  if distutils.version.LooseVersion(tf.__version__) >= "2.2":
+    base_dirs = [
+        base_dir,
+        path.dirname(tensorboard.__file__),
+        path.dirname(tensorflow_estimator.__file__),
+    ]
+  else:
+    base_dirs = [
+        path.normpath(path.join(base_dir, "../tensorflow_core")),
+        path.dirname(tensorboard.__file__),
+        path.dirname(tensorflow_estimator.__file__),
+    ]
+
+  code_url_prefixes = (
+      code_url_prefix,
+      "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
+      "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
+  )
+
+  return base_dirs, code_url_prefixes

tensorflow/tools/docs/generate2.py

@@ -30,7 +30,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from os import path
+import pathlib
 import textwrap
 
 from absl import app
@@ -42,12 +42,13 @@ 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
 
-import tensorboard
-import tensorflow_estimator
 from tensorflow.python.framework import ops
 from tensorflow.python.util import tf_export
 from tensorflow.python.util import tf_inspect
 
+# Caution: the google and oss versions of this import are different.
+import base_dir
+
 # `tf` has an `__all__` that doesn't list important things like `keras`.
 # The doc generator recognizes `__all__` as the list of public symbols.
 # So patch `tf.__all__` to list everything.
@@ -202,22 +203,8 @@ def build_docs(output_dir, code_url_prefix, search_hints=True):
   except AttributeError:
     pass
 
-  base_dir = path.normpath(path.join(tf.__file__, "../.."))
-
-  base_dirs = (
-      path.join(base_dir, "tensorflow_core"),
-      # External packages base directories
-      path.dirname(tensorboard.__file__),
-      path.dirname(tensorflow_estimator.__file__),
-  )
-
-  code_url_prefixes = (
-      code_url_prefix,
-      # External packages source repositories,
-      "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
-      "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
-  )
-
+  base_dirs, code_url_prefixes = base_dir.get_base_dirs_and_prefixes(
+      code_url_prefix)
   doc_generator = generate_lib.DocGenerator(
       root_title="TensorFlow 2",
       py_modules=[("tf", tf)],
@@ -230,6 +217,38 @@ def build_docs(output_dir, code_url_prefix, search_hints=True):
 
   doc_generator.build(output_dir)
 
+  out_path = pathlib.Path(output_dir)
+  num_files = len(list(out_path.rglob("*")))
+  if num_files < 2500:
+    raise ValueError("The TensorFlow api should be more than 2500 files"
+                     "(found {}).".format(num_files))
+  expected_path_contents = {
+      "tf/summary/audio.md":
+          "tensorboard/plugins/audio/summary_v2.py",
+      "tf/estimator/DNNClassifier.md":
+          "tensorflow_estimator/python/estimator/canned/dnn.py",
+      "tf/nn/sigmoid_cross_entropy_with_logits.md":
+          "python/ops/nn_impl.py",
+      "tf/keras/Model.md":
+          "tensorflow/python/keras/engine/training.py",
+      "tf/compat/v1/gradients.md":
+          "tensorflow/python/ops/gradients_impl.py",
+  }
+
+  all_passed = True
+  error_msg_parts = [
+      'Some "view source" links seem to be broken, please check:'
+  ]
+
+  for (rel_path, contents) in expected_path_contents.items():
+    path = out_path / rel_path
+    if contents not in path.read_text():
+      all_passed = False
+      error_msg_parts.append("  " + str(path))
+
+  if not all_passed:
+    raise ValueError("\n".join(error_msg_parts))
+
 
 def main(argv):
   del argv