xref: /aosp_15_r20/external/tensorflow/tensorflow/python/util/all_util.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
16"""Generate __all__ from a module docstring."""
17import re as _re
18import sys as _sys
19
20from tensorflow.python.util import tf_inspect as _tf_inspect
21
22
23_reference_pattern = _re.compile(r'^@@(\w+)$', flags=_re.MULTILINE)
24
25
26def make_all(module_name, doc_string_modules=None):
27  """Generates `__all__` from the docstring of one or more modules.
28
29  Usage: `make_all(__name__)` or
30  `make_all(__name__, [sys.modules(__name__), other_module])`. The doc string
31  modules must each a docstring, and `__all__` will contain all symbols with
32  `@@` references, where that symbol currently exists in the module named
33  `module_name`.
34
35  Args:
36    module_name: The name of the module (usually `__name__`).
37    doc_string_modules: a list of modules from which to take docstring.
38    If None, then a list containing only the module named `module_name` is used.
39
40  Returns:
41    A list suitable for use as `__all__`.
42  """
43  if doc_string_modules is None:
44    doc_string_modules = [_sys.modules[module_name]]
45  cur_members = set(
46      name for name, _ in _tf_inspect.getmembers(_sys.modules[module_name]))
47
48  results = set()
49  for doc_module in doc_string_modules:
50    results.update([m.group(1)
51                    for m in _reference_pattern.finditer(doc_module.__doc__)
52                    if m.group(1) in cur_members])
53  return list(results)
54
55# Hidden attributes are attributes that have been hidden by
56# `remove_undocumented`. They can be re-instated by `reveal_undocumented`.
57# This maps symbol names to a tuple, containing:
58#   (module object, attribute value)
59_HIDDEN_ATTRIBUTES = {}
60
61
62def reveal_undocumented(symbol_name, target_module=None):
63  """Reveals a symbol that was previously removed by `remove_undocumented`.
64
65  This should be used by tensorflow internal tests only. It explicitly
66  defeats the encapsulation afforded by `remove_undocumented`.
67
68  It throws an exception when the symbol was not hidden in the first place.
69
70  Args:
71    symbol_name: a string representing the full absolute path of the symbol.
72    target_module: if specified, the module in which to restore the symbol.
73  """
74  if symbol_name not in _HIDDEN_ATTRIBUTES:
75    raise LookupError('Symbol %s is not a hidden symbol' % symbol_name)
76  symbol_basename = symbol_name.split('.')[-1]
77  (original_module, attr_value) = _HIDDEN_ATTRIBUTES[symbol_name]
78  if not target_module: target_module = original_module
79  setattr(target_module, symbol_basename, attr_value)
80
81
82def remove_undocumented(module_name, allowed_exception_list=None,
83                        doc_string_modules=None):
84  """Removes symbols in a module that are not referenced by a docstring.
85
86  Args:
87    module_name: the name of the module (usually `__name__`).
88    allowed_exception_list: a list of names that should not be removed.
89    doc_string_modules: a list of modules from which to take the docstrings.
90    If None, then a list containing only the module named `module_name` is used.
91
92    Furthermore, if a symbol previously added with `add_to_global_allowlist`,
93    then it will always be allowed. This is useful for internal tests.
94
95  Returns:
96    None
97  """
98  current_symbols = set(dir(_sys.modules[module_name]))
99  should_have = make_all(module_name, doc_string_modules)
100  should_have += allowed_exception_list or []
101  extra_symbols = current_symbols - set(should_have)
102  target_module = _sys.modules[module_name]
103  for extra_symbol in extra_symbols:
104    # Skip over __file__, etc. Also preserves internal symbols.
105    if extra_symbol.startswith('_'): continue
106    fully_qualified_name = module_name + '.' + extra_symbol
107    _HIDDEN_ATTRIBUTES[fully_qualified_name] = (target_module,
108                                                getattr(target_module,
109                                                        extra_symbol))
110    delattr(target_module, extra_symbol)
111
112
113__all__ = [
114    'make_all',
115    'remove_undocumented',
116    'reveal_undocumented',
117]
118