aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/tools
diff options
context:
space:
mode:
authorGravatar Michael Case <mikecase@google.com>2018-08-18 19:33:13 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-08-18 19:36:33 -0700
commit51100a8de57ef53e36a8a9f5a9829cbd33fbed04 (patch)
tree50fcca9d7b912adf82e047f9cca22bfa7f8f21cb /tensorflow/python/tools
parent424ce99f4357515f0f9e57b7a087194cf46134d6 (diff)
Internal Change.
PiperOrigin-RevId: 209299599
Diffstat (limited to 'tensorflow/python/tools')
-rw-r--r--tensorflow/python/tools/BUILD6
-rw-r--r--tensorflow/python/tools/component_api_helper.py85
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()