Introduce TFDecorator, a base class for Python TensorFlow decorators. Provides basic introspection and "unwrap" services, allowing tooling code to fully 'understand' the wrapped object.

Change: 153854044

2 files changed, 10 insertions, 8 deletions

diff --git a/tensorflow/tools/common/public_api.py b/tensorflow/tools/common/public_api.py
index 44c6f24e56..837f11f690 100644
--- a/tensorflow/tools/common/public_api.py
+++ b/tensorflow/tools/common/public_api.py
@@ -18,9 +18,10 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 import re
 
+from tensorflow.python.util import tf_inspect
+
 
 class PublicAPIVisitor(object):
   """Visitor to use with `traverse` to visit exactly the public TF API."""
@@ -93,7 +94,7 @@ class PublicAPIVisitor(object):
     """Visitor interface, see `traverse` for details."""
 
     # Avoid long waits in cases of pretty unambiguous failure.
-    if inspect.ismodule(parent) and len(path.split('.')) > 10:
+    if tf_inspect.ismodule(parent) and len(path.split('.')) > 10:
       raise RuntimeError('Modules nested too deep:\n%s\n\nThis is likely a '
                          'problem with an accidental public import.' % path)
 
diff --git a/tensorflow/tools/common/traverse.py b/tensorflow/tools/common/traverse.py
index 443838d968..9607f80686 100644
--- a/tensorflow/tools/common/traverse.py
+++ b/tensorflow/tools/common/traverse.py
@@ -18,9 +18,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import inspect
 import sys
 
+from tensorflow.python.util import tf_inspect
 
 __all__ = ['traverse']
 
@@ -29,11 +29,11 @@ def _traverse_internal(root, visit, stack, path):
   """Internal helper for traverse."""
 
   # Only traverse modules and classes
-  if not inspect.isclass(root) and not inspect.ismodule(root):
+  if not tf_inspect.isclass(root) and not tf_inspect.ismodule(root):
     return
 
   try:
-    children = inspect.getmembers(root)
+    children = tf_inspect.getmembers(root)
   except ImportError:
     # On some Python installations, some modules do not support enumerating
     # members (six in particular), leading to import errors.
@@ -43,7 +43,8 @@ def _traverse_internal(root, visit, stack, path):
   visit(path, root, children)
   for name, child in children:
     # Do not descend into built-in modules
-    if inspect.ismodule(child) and child.__name__ in sys.builtin_module_names:
+    if tf_inspect.ismodule(
+        child) and child.__name__ in sys.builtin_module_names:
       continue
 
     # Break cycles
@@ -72,8 +73,8 @@ def traverse(root, visit):
   never descends into built-in modules.
 
   `children`, a list of `(name, object)` pairs are determined by
-  `inspect.getmembers`. To avoid visiting parts of the tree, `children` can be
-  modified in place, using `del` or slice assignment.
+  `tf_inspect.getmembers`. To avoid visiting parts of the tree, `children` can
+  be modified in place, using `del` or slice assignment.
 
   Cycles (determined by reference equality, `is`) stop the traversal. A stack of
   objects is kept to find cycles. Objects forming cycles may appear in