Make a silencing context manager for deprecations and use it.

PiperOrigin-RevId: 156870702

2 files changed, 81 insertions, 44 deletions

diff --git a/tensorflow/python/util/deprecation.py b/tensorflow/python/util/deprecation.py
index 1e1599afb4..c2de723bfb 100644
--- a/tensorflow/python/util/deprecation.py
+++ b/tensorflow/python/util/deprecation.py
@@ -24,10 +24,15 @@ import re
 
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import decorator_utils
+from tensorflow.python.util import tf_contextlib
 from tensorflow.python.util import tf_decorator
 from tensorflow.python.util import tf_inspect
 
 
+# Allow deprecation warnings to be silenced temporarily with a context manager.
+_PRINT_DEPRECATION_WARNINGS = True
+
+
 def _add_deprecated_function_notice_to_docstring(doc, date, instructions):
   """Adds a deprecation notice to a docstring for deprecated functions."""
   return decorator_utils.add_notice_to_docstring(
@@ -111,14 +116,15 @@ def deprecated(date, instructions):
     """Deprecation wrapper."""
     decorator_utils.validate_callable(func, 'deprecated')
     @functools.wraps(func)
-    def new_func(*args, **kwargs):
-      logging.warning(
-          'From %s: %s (from %s) is deprecated and will be removed %s.\n'
-          'Instructions for updating:\n%s',
-          _call_location(), decorator_utils.get_qualified_name(func),
-          func.__module__,
-          'in a future version' if date is None else ('after %s' % date),
-          instructions)
+    def new_func(*args, **kwargs):  # pylint: disable=missing-docstring
+      if _PRINT_DEPRECATION_WARNINGS:
+        logging.warning(
+            'From %s: %s (from %s) is deprecated and will be removed %s.\n'
+            'Instructions for updating:\n%s',
+            _call_location(), decorator_utils.get_qualified_name(func),
+            func.__module__,
+            'in a future version' if date is None else ('after %s' % date),
+            instructions)
       return func(*args, **kwargs)
     return tf_decorator.make_decorator(
         func, new_func, 'deprecated',
@@ -261,31 +267,32 @@ def deprecated_args(date, instructions, *deprecated_arg_names_or_tuples):
     @functools.wraps(func)
     def new_func(*args, **kwargs):
       """Deprecation wrapper."""
-      invalid_args = []
-      named_args = tf_inspect.getcallargs(func, *args, **kwargs)
-      for arg_name, spec in iter(deprecated_positions.items()):
-        if (spec.position < len(args) and
-            not (spec.has_ok_value and
-                 _same_value(named_args[arg_name], spec.ok_value))):
-          invalid_args.append(arg_name)
-      if is_varargs_deprecated and len(args) > len(arg_spec.args):
-        invalid_args.append(arg_spec.varargs)
-      if is_kwargs_deprecated and kwargs:
-        invalid_args.append(arg_spec.keywords)
-      for arg_name in deprecated_arg_names:
-        if (arg_name in kwargs and
-            not (deprecated_positions[arg_name].has_ok_value and
-                 _same_value(named_args[arg_name],
-                             deprecated_positions[arg_name].ok_value))):
-          invalid_args.append(arg_name)
-      for arg_name in invalid_args:
-        logging.warning(
-            'From %s: calling %s (from %s) with %s is deprecated and will '
-            'be removed %s.\nInstructions for updating:\n%s',
-            _call_location(), decorator_utils.get_qualified_name(func),
-            func.__module__, arg_name,
-            'in a future version' if date is None else ('after %s' % date),
-            instructions)
+      if _PRINT_DEPRECATION_WARNINGS:
+        invalid_args = []
+        named_args = tf_inspect.getcallargs(func, *args, **kwargs)
+        for arg_name, spec in iter(deprecated_positions.items()):
+          if (spec.position < len(args) and
+              not (spec.has_ok_value and
+                   _same_value(named_args[arg_name], spec.ok_value))):
+            invalid_args.append(arg_name)
+        if is_varargs_deprecated and len(args) > len(arg_spec.args):
+          invalid_args.append(arg_spec.varargs)
+        if is_kwargs_deprecated and kwargs:
+          invalid_args.append(arg_spec.keywords)
+        for arg_name in deprecated_arg_names:
+          if (arg_name in kwargs and
+              not (deprecated_positions[arg_name].has_ok_value and
+                   _same_value(named_args[arg_name],
+                               deprecated_positions[arg_name].ok_value))):
+            invalid_args.append(arg_name)
+        for arg_name in invalid_args:
+          logging.warning(
+              'From %s: calling %s (from %s) with %s is deprecated and will '
+              'be removed %s.\nInstructions for updating:\n%s',
+              _call_location(), decorator_utils.get_qualified_name(func),
+              func.__module__, arg_name,
+              'in a future version' if date is None else ('after %s' % date),
+              instructions)
       return func(*args, **kwargs)
     return tf_decorator.make_decorator(func, new_func, 'deprecated',
                                        _add_deprecated_arg_notice_to_docstring(
@@ -334,16 +341,17 @@ def deprecated_arg_values(date, instructions, **deprecated_kwargs):
     @functools.wraps(func)
     def new_func(*args, **kwargs):
       """Deprecation wrapper."""
-      named_args = tf_inspect.getcallargs(func, *args, **kwargs)
-      for arg_name, arg_value in deprecated_kwargs.items():
-        if arg_name in named_args and named_args[arg_name] == arg_value:
-          logging.warning(
-              'From %s: calling %s (from %s) with %s=%s is deprecated and will '
-              'be removed %s.\nInstructions for updating:\n%s',
-              _call_location(), decorator_utils.get_qualified_name(func),
-              func.__module__, arg_name, arg_value,
-              'in a future version' if date is None else ('after %s' % date),
-              instructions)
+      if _PRINT_DEPRECATION_WARNINGS:
+        named_args = tf_inspect.getcallargs(func, *args, **kwargs)
+        for arg_name, arg_value in deprecated_kwargs.items():
+          if arg_name in named_args and named_args[arg_name] == arg_value:
+            logging.warning(
+                'From %s: calling %s (from %s) with %s=%s is deprecated and '
+                'will be removed %s.\nInstructions for updating:\n%s',
+                _call_location(), decorator_utils.get_qualified_name(func),
+                func.__module__, arg_name, arg_value,
+                'in a future version' if date is None else ('after %s' % date),
+                instructions)
       return func(*args, **kwargs)
     return tf_decorator.make_decorator(func, new_func, 'deprecated',
                                        _add_deprecated_arg_notice_to_docstring(
@@ -375,3 +383,13 @@ def deprecated_argument_lookup(new_name, new_value, old_name, old_value):
 def rewrite_argument_docstring(old_doc, old_argument, new_argument):
   return old_doc.replace('`%s`' % old_argument, '`%s`' % new_argument).replace(
       '%s:' % old_argument, '%s:' % new_argument)
+
+
+@tf_contextlib.contextmanager
+def silence():
+  """Temporarily silence deprecation warnings."""
+  global _PRINT_DEPRECATION_WARNINGS
+  print_deprecation_warnings = _PRINT_DEPRECATION_WARNINGS
+  _PRINT_DEPRECATION_WARNINGS = False
+  yield
+  _PRINT_DEPRECATION_WARNINGS = print_deprecation_warnings
diff --git a/tensorflow/python/util/deprecation_test.py b/tensorflow/python/util/deprecation_test.py
index cce0bb1b4e..e2d9a594a3 100644
--- a/tensorflow/python/util/deprecation_test.py
+++ b/tensorflow/python/util/deprecation_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""tensor_util tests."""
+"""Deprecation tests."""
 
 # pylint: disable=unused-import
 from __future__ import absolute_import
@@ -26,6 +26,25 @@ from tensorflow.python.util import deprecation
 
 class DeprecationTest(test.TestCase):
 
+  @test.mock.patch.object(logging, "warning", autospec=True)
+  def test_silence(self, mock_warning):
+    date = "2016-07-04"
+    instructions = "This is how you update..."
+
+    @deprecation.deprecated(date, instructions)
+    def _fn():
+      pass
+
+    _fn()
+    self.assertEqual(1, mock_warning.call_count)
+
+    with deprecation.silence():
+      _fn()
+    self.assertEqual(1, mock_warning.call_count)
+
+    _fn()
+    self.assertEqual(2, mock_warning.call_count)
+
   def _assert_subset(self, expected_subset, actual_set):
     self.assertTrue(
         actual_set.issuperset(expected_subset),