xref: /aosp_15_r20/external/tensorflow/tensorflow/python/util/module_wrapper.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2021 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"""Provides wrapper for TensorFlow modules."""
16
17import importlib
18
19from tensorflow.python.eager import monitoring
20from tensorflow.python.platform import tf_logging as logging
21from tensorflow.python.util import fast_module_type
22from tensorflow.python.util import tf_decorator
23from tensorflow.python.util import tf_inspect
24from tensorflow.tools.compatibility import all_renames_v2
25
26FastModuleType = fast_module_type.get_fast_module_type_class()
27_PER_MODULE_WARNING_LIMIT = 1
28compat_v1_usage_gauge = monitoring.BoolGauge('/tensorflow/api/compat/v1',
29                                             'compat.v1 usage')
30
31
32def get_rename_v2(name):
33  if name not in all_renames_v2.symbol_renames:
34    return None
35  return all_renames_v2.symbol_renames[name]
36
37
38def _call_location():
39  """Extracts the caller filename and line number as a string.
40
41  Returns:
42    A string describing the caller source location.
43  """
44  frame = tf_inspect.currentframe()
45  assert frame.f_back.f_code.co_name == '_tfmw_add_deprecation_warning', (
46      'This function should be called directly from '
47      '_tfmw_add_deprecation_warning, as the caller is identified '
48      'heuristically by chopping off the top stack frames.')
49
50  # We want to get stack frame 3 frames up from current frame,
51  # i.e. above __getattr__, _tfmw_add_deprecation_warning,
52  # and _call_location calls.
53  for _ in range(3):
54    parent = frame.f_back
55    if parent is None:
56      break
57    frame = parent
58  return '{}:{}'.format(frame.f_code.co_filename, frame.f_lineno)
59
60
61def contains_deprecation_decorator(decorators):
62  return any(d.decorator_name == 'deprecated' for d in decorators)
63
64
65def has_deprecation_decorator(symbol):
66  """Checks if given object has a deprecation decorator.
67
68  We check if deprecation decorator is in decorators as well as
69  whether symbol is a class whose __init__ method has a deprecation
70  decorator.
71  Args:
72    symbol: Python object.
73
74  Returns:
75    True if symbol has deprecation decorator.
76  """
77  decorators, symbol = tf_decorator.unwrap(symbol)
78  if contains_deprecation_decorator(decorators):
79    return True
80  if tf_inspect.isfunction(symbol):
81    return False
82  if not tf_inspect.isclass(symbol):
83    return False
84  if not hasattr(symbol, '__init__'):
85    return False
86  init_decorators, _ = tf_decorator.unwrap(symbol.__init__)
87  return contains_deprecation_decorator(init_decorators)
88
89
90class TFModuleWrapper(FastModuleType):
91  """Wrapper for TF modules to support deprecation messages and lazyloading."""
92  # Ensures that compat.v1 API usage is recorded at most once
93  compat_v1_usage_recorded = False
94
95  def __init__(
96      self,
97      wrapped,
98      module_name,
99      public_apis=None,
100      deprecation=True,
101      has_lite=False):
102    super(TFModuleWrapper, self).__init__(wrapped.__name__)
103    FastModuleType.set_getattr_callback(self, TFModuleWrapper._getattr)
104    FastModuleType.set_getattribute_callback(self,
105                                             TFModuleWrapper._getattribute)
106    self.__dict__.update(wrapped.__dict__)
107    # Prefix all local attributes with _tfmw_ so that we can
108    # handle them differently in attribute access methods.
109    self._tfmw_wrapped_module = wrapped
110    self._tfmw_module_name = module_name
111    self._tfmw_public_apis = public_apis
112    self._tfmw_print_deprecation_warnings = deprecation
113    self._tfmw_has_lite = has_lite
114    self._tfmw_is_compat_v1 = (wrapped.__name__.endswith('.compat.v1'))
115    # Set __all__ so that import * work for lazy loaded modules
116    if self._tfmw_public_apis:
117      self._tfmw_wrapped_module.__all__ = list(self._tfmw_public_apis.keys())
118      self.__all__ = list(self._tfmw_public_apis.keys())
119    else:
120      if hasattr(self._tfmw_wrapped_module, '__all__'):
121        self.__all__ = self._tfmw_wrapped_module.__all__
122      else:
123        self._tfmw_wrapped_module.__all__ = [
124            attr for attr in dir(self._tfmw_wrapped_module)
125            if not attr.startswith('_')
126        ]
127        self.__all__ = self._tfmw_wrapped_module.__all__
128
129    # names we already checked for deprecation
130    self._tfmw_deprecated_checked = set()
131    self._tfmw_warning_count = 0
132
133  def _tfmw_add_deprecation_warning(self, name, attr):
134    """Print deprecation warning for attr with given name if necessary."""
135    if (self._tfmw_warning_count < _PER_MODULE_WARNING_LIMIT and
136        name not in self._tfmw_deprecated_checked):
137
138      self._tfmw_deprecated_checked.add(name)
139
140      if self._tfmw_module_name:
141        full_name = 'tf.%s.%s' % (self._tfmw_module_name, name)
142      else:
143        full_name = 'tf.%s' % name
144      rename = get_rename_v2(full_name)
145      if rename and not has_deprecation_decorator(attr):
146        call_location = _call_location()
147        # skip locations in Python source
148        if not call_location.startswith('<'):
149          logging.warning(
150              'From %s: The name %s is deprecated. Please use %s instead.\n',
151              _call_location(), full_name, rename)
152          self._tfmw_warning_count += 1
153          return True
154    return False
155
156  def _tfmw_import_module(self, name):
157    """Lazily loading the modules."""
158    # We ignore 'app' because it is accessed in __init__.py of tf.compat.v1.
159    # That way, if a user only imports tensorflow.compat.v1, it is not
160    # considered v1 API usage.
161    if (self._tfmw_is_compat_v1 and name != 'app' and
162        not TFModuleWrapper.compat_v1_usage_recorded):
163      TFModuleWrapper.compat_v1_usage_recorded = True
164      compat_v1_usage_gauge.get_cell().set(True)
165
166    symbol_loc_info = self._tfmw_public_apis[name]
167    if symbol_loc_info[0]:
168      module = importlib.import_module(symbol_loc_info[0])
169      attr = getattr(module, symbol_loc_info[1])
170    else:
171      attr = importlib.import_module(symbol_loc_info[1])
172    setattr(self._tfmw_wrapped_module, name, attr)
173    self.__dict__[name] = attr
174    # Cache the pair
175    self._fastdict_insert(name, attr)
176    return attr
177
178  def _getattribute(self, name):
179    # pylint: disable=g-doc-return-or-yield,g-doc-args
180    """Imports and caches pre-defined API.
181
182    Warns if necessary.
183
184    This method is a replacement for __getattribute__(). It will be added into
185    the extended python module as a callback to reduce API overhead.
186    """
187    # Avoid infinite recursions
188    func__fastdict_insert = object.__getattribute__(self, '_fastdict_insert')
189
190    # Make sure we do not import from tensorflow/lite/__init__.py
191    if name == 'lite':
192      if self._tfmw_has_lite:
193        attr = self._tfmw_import_module(name)
194        setattr(self._tfmw_wrapped_module, 'lite', attr)
195        func__fastdict_insert(name, attr)
196        return attr
197  # Placeholder for Google-internal contrib error
198
199    attr = object.__getattribute__(self, name)
200
201    # Return and cache dunders and our own members.
202    # This is necessary to guarantee successful construction.
203    # In addition, all the accessed attributes used during the construction must
204    # begin with "__" or "_tfmw" or "_fastdict_".
205    if name.startswith('__') or name.startswith('_tfmw_') or name.startswith(
206        '_fastdict_'):
207      func__fastdict_insert(name, attr)
208      return attr
209
210    # Print deprecations, only cache functions after deprecation warnings have
211    # stopped.
212    if not (self._tfmw_print_deprecation_warnings and
213            self._tfmw_add_deprecation_warning(name, attr)):
214      func__fastdict_insert(name, attr)
215
216    return attr
217
218  def _getattr(self, name):
219    # pylint: disable=g-doc-return-or-yield,g-doc-args
220    """Imports and caches pre-defined API.
221
222    Warns if necessary.
223
224    This method is a replacement for __getattr__(). It will be added into the
225    extended python module as a callback to reduce API overhead. Instead of
226    relying on implicit AttributeError handling, this added callback function
227    will
228    be called explicitly from the extended C API if the default attribute lookup
229    fails.
230    """
231    try:
232      attr = getattr(self._tfmw_wrapped_module, name)
233    except AttributeError:
234    # Placeholder for Google-internal contrib error
235
236      if not self._tfmw_public_apis:
237        raise
238      if name not in self._tfmw_public_apis:
239        raise
240      attr = self._tfmw_import_module(name)
241
242    if self._tfmw_print_deprecation_warnings:
243      self._tfmw_add_deprecation_warning(name, attr)
244    return attr
245
246  def __setattr__(self, arg, val):
247    if not arg.startswith('_tfmw_'):
248      setattr(self._tfmw_wrapped_module, arg, val)
249      self.__dict__[arg] = val
250      if arg not in self.__all__ and arg != '__all__':
251        self.__all__.append(arg)
252      # Update the cache
253      if self._fastdict_key_in(arg):
254        self._fastdict_insert(arg, val)
255    super(TFModuleWrapper, self).__setattr__(arg, val)
256
257  def __dir__(self):
258    if self._tfmw_public_apis:
259      return list(
260          set(self._tfmw_public_apis.keys()).union(
261              set([
262                  attr for attr in dir(self._tfmw_wrapped_module)
263                  if not attr.startswith('_')
264              ])))
265    else:
266      return dir(self._tfmw_wrapped_module)
267
268  def __delattr__(self, name):
269    if name.startswith('_tfmw_'):
270      super(TFModuleWrapper, self).__delattr__(name)
271    else:
272      delattr(self._tfmw_wrapped_module, name)
273
274  def __repr__(self):
275    return self._tfmw_wrapped_module.__repr__()
276
277  def __reduce__(self):
278    return importlib.import_module, (self.__name__,)
279