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