• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Visitor restricting traversal to only the public tensorflow API."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import re
22
23from tensorflow.python.util import tf_inspect
24
25
26class PublicAPIVisitor(object):
27  """Visitor to use with `traverse` to visit exactly the public TF API."""
28
29  def __init__(self, visitor):
30    """Constructor.
31
32    `visitor` should be a callable suitable as a visitor for `traverse`. It will
33    be called only for members of the public TensorFlow API.
34
35    Args:
36      visitor: A visitor to call for the public API.
37    """
38    self._visitor = visitor
39    self._root_name = 'tf'
40
41    # Modules/classes we want to suppress entirely.
42    self._private_map = {
43        # Some implementations have this internal module that we shouldn't
44        # expose.
45        'tf.flags': ['cpp_flags'],
46    }
47
48    # Modules/classes we do not want to descend into if we hit them. Usually,
49    # system modules exposed through platforms for compatibility reasons.
50    # Each entry maps a module path to a name to ignore in traversal.
51    self._do_not_descend_map = {
52        'tf': [
53            'compiler',
54            'core',
55            'examples',
56            'flags',  # Don't add flags
57            # TODO(drpng): This can be removed once sealed off.
58            'platform',
59            # TODO(drpng): This can be removed once sealed.
60            'pywrap_tensorflow',
61            # TODO(drpng): This can be removed once sealed.
62            'user_ops',
63            'python',
64            'tools',
65            'tensorboard',
66        ],
67
68        ## Everything below here is legitimate.
69        # It'll stay, but it's not officially part of the API.
70        'tf.app': ['flags'],
71        # Imported for compatibility between py2/3.
72        'tf.test': ['mock'],
73        # Externalized modules of the Keras API.
74        'tf.keras': ['applications', 'preprocessing']
75    }
76
77  @property
78  def private_map(self):
79    """A map from parents to symbols that should not be included at all.
80
81    This map can be edited, but it should not be edited once traversal has
82    begun.
83
84    Returns:
85      The map marking symbols to not include.
86    """
87    return self._private_map
88
89  @property
90  def do_not_descend_map(self):
91    """A map from parents to symbols that should not be descended into.
92
93    This map can be edited, but it should not be edited once traversal has
94    begun.
95
96    Returns:
97      The map marking symbols to not explore.
98    """
99    return self._do_not_descend_map
100
101  def set_root_name(self, root_name):
102    """Override the default root name of 'tf'."""
103    self._root_name = root_name
104
105  def _is_private(self, path, name, obj=None):
106    """Return whether a name is private."""
107    # TODO(wicke): Find out what names to exclude.
108    del obj  # Unused.
109    return ((path in self._private_map and
110             name in self._private_map[path]) or
111            (name.startswith('_') and not re.match('__.*__$', name) or
112             name in ['__base__', '__class__']))
113
114  def _do_not_descend(self, path, name):
115    """Safely queries if a specific fully qualified name should be excluded."""
116    return (path in self._do_not_descend_map and
117            name in self._do_not_descend_map[path])
118
119  def __call__(self, path, parent, children):
120    """Visitor interface, see `traverse` for details."""
121
122    # Avoid long waits in cases of pretty unambiguous failure.
123    if tf_inspect.ismodule(parent) and len(path.split('.')) > 10:
124      raise RuntimeError('Modules nested too deep:\n%s.%s\n\nThis is likely a '
125                         'problem with an accidental public import.' %
126                         (self._root_name, path))
127
128    # Includes self._root_name
129    full_path = '.'.join([self._root_name, path]) if path else self._root_name
130
131    # Remove things that are not visible.
132    for name, child in list(children):
133      if self._is_private(full_path, name, child):
134        children.remove((name, child))
135
136    self._visitor(path, parent, children)
137
138    # Remove things that are visible, but which should not be descended into.
139    for name, child in list(children):
140      if self._do_not_descend(full_path, name):
141        children.remove((name, child))
142