Internal Change.

PiperOrigin-RevId: 209299599

2 files changed, 91 insertions, 0 deletions

diff --git a/tensorflow/python/tools/BUILD b/tensorflow/python/tools/BUILD
index 222f856511..01d43e09d1 100644
--- a/tensorflow/python/tools/BUILD
+++ b/tensorflow/python/tools/BUILD
@@ -114,6 +114,12 @@ py_library(
     ],
 )
 
+py_library(
+    name = "component_api_helper",
+    srcs = ["component_api_helper.py"],
+    srcs_version = "PY2AND3",
+)
+
 py_binary(
     name = "strip_unused",
     srcs = ["strip_unused.py"],
diff --git a/tensorflow/python/tools/component_api_helper.py b/tensorflow/python/tools/component_api_helper.py
new file mode 100644
index 0000000000..988ecc61f0
--- /dev/null
+++ b/tensorflow/python/tools/component_api_helper.py
@@ -0,0 +1,85 @@
+# 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.
+# ==============================================================================
+"""Helper functions to help integrate TensorFlow components into TF API.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import importlib
+import os
+
+
+def package_hook(parent_package_str, child_package_str, error_msg=None):
+  """Used to hook in an external package into the TensorFlow namespace.
+
+  Example usage:
+  ### tensorflow/__init__.py
+  from tensorflow.python.tools import component_api_helper
+  component_api_helper.package_hook(
+      'tensorflow', 'tensorflow_estimator.python')
+  component_api_helper(
+      'tensorflow.contrib', 'tensorflow_estimator.contrib.python')
+  del component_api_helper
+
+  TODO(mikecase): This function has a minor issue, where if the child package
+  does not exist alone in its directory, sibling packages to it will also be
+  accessible from the parent. This is because we just add
+  `child_pkg.__file__/..` to the subpackage search path. This should not be
+  a big issue because of how our API generation scripts work (the child package
+  we are hooking up should always be alone). But there might be a better way
+  of doing this.
+
+  Args:
+    parent_package_str: Parent package name as a string such as 'tensorflow' or
+      'tensorflow.contrib'. This will become the parent package for the
+      component package being hooked in.
+    child_package_str: Child package name as a string such as
+      'tensorflow_estimator.python'. This package will be added as a subpackage
+      of the parent.
+    error_msg: Message to print if child package cannot be found.
+  """
+  parent_pkg = importlib.import_module(parent_package_str)
+  try:
+    child_pkg = importlib.import_module(child_package_str)
+  except ImportError:
+    if error_msg:
+      print(error_msg)
+    return
+
+  def set_child_as_subpackage():
+    """Sets child package as a subpackage of parent package.
+
+    Will allow the following import statement to work.
+    >>> import parent.child
+    """
+    child_pkg_path = [os.path.join(os.path.dirname(child_pkg.__file__), "..")]
+    try:
+      parent_pkg.__path__ += child_pkg_path
+    except AttributeError:
+      parent_pkg.__path__ = child_pkg_path
+
+  def set_child_as_attr():
+    """Sets child package as a attr of the parent package.
+
+    Will allow for the following.
+    >>> import parent
+    >>> parent.child
+    """
+    child_pkg_attr_name = child_pkg.__name__.split(".")[-1]
+    setattr(parent_pkg, child_pkg_attr_name, child_pkg)
+
+  set_child_as_subpackage()
+  set_child_as_attr()