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"""Traversing Python modules and classes."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import sys
22
23from tensorflow.python.util import tf_inspect
24
25__all__ = ['traverse']
26
27
28def _traverse_internal(root, visit, stack, path):
29  """Internal helper for traverse."""
30
31  # Only traverse modules and classes
32  if not tf_inspect.isclass(root) and not tf_inspect.ismodule(root):
33    return
34
35  try:
36    children = tf_inspect.getmembers(root)
37  except ImportError:
38    # On some Python installations, some modules do not support enumerating
39    # members (six in particular), leading to import errors.
40    children = []
41
42  new_stack = stack + [root]
43  visit(path, root, children)
44  for name, child in children:
45    # Do not descend into built-in modules
46    if tf_inspect.ismodule(
47        child) and child.__name__ in sys.builtin_module_names:
48      continue
49
50    # Break cycles
51    if any(child is item for item in new_stack):  # `in`, but using `is`
52      continue
53
54    child_path = path + '.' + name if path else name
55    _traverse_internal(child, visit, new_stack, child_path)
56
57
58def traverse(root, visit):
59  """Recursively enumerate all members of `root`.
60
61  Similar to the Python library function `os.path.walk`.
62
63  Traverses the tree of Python objects starting with `root`, depth first.
64  Parent-child relationships in the tree are defined by membership in modules or
65  classes. The function `visit` is called with arguments
66  `(path, parent, children)` for each module or class `parent` found in the tree
67  of python objects starting with `root`. `path` is a string containing the name
68  with which `parent` is reachable from the current context. For example, if
69  `root` is a local class called `X` which contains a class `Y`, `visit` will be
70  called with `('Y', X.Y, children)`).
71
72  If `root` is not a module or class, `visit` is never called. `traverse`
73  never descends into built-in modules.
74
75  `children`, a list of `(name, object)` pairs are determined by
76  `tf_inspect.getmembers`. To avoid visiting parts of the tree, `children` can
77  be modified in place, using `del` or slice assignment.
78
79  Cycles (determined by reference equality, `is`) stop the traversal. A stack of
80  objects is kept to find cycles. Objects forming cycles may appear in
81  `children`, but `visit` will not be called with any object as `parent` which
82  is already in the stack.
83
84  Traversing system modules can take a long time, it is advisable to pass a
85  `visit` callable which blacklists such modules.
86
87  Args:
88    root: A python object with which to start the traversal.
89    visit: A function taking arguments `(path, parent, children)`. Will be
90      called for each object found in the traversal.
91  """
92  _traverse_internal(root, visit, [], '')
93