assert_deprecated improvements (#3919)

Deprecates the `allow_multiple_warnings` parameter for explicit `count`. 

Context: this is split out from #3917.

===== Release note ======

- `cirq.testing.assert_deprecate` parameter `allow_multiple_warnings` is now deprecated and will be removed in Cirq v0.12. Please use instead `count`. For `allow_multiple_warnings` `True`, you can just use `count=None`, for `False`, `count=1` which is the default value.

======================

cirq/_compat.py

@@ -17,8 +17,8 @@
 import os
 import re
 import warnings
-from typing import Any, Callable, Optional, Dict, Tuple, Type
 from types import ModuleType
+from typing import Any, Callable, Optional, Dict, Tuple, Type
 
 import numpy as np
 import pandas as pd
@@ -96,7 +96,7 @@ def _warn_or_error(msg):
     warnings.warn(
         msg,
         DeprecationWarning,
-        stacklevel=2,
+        stacklevel=3,
     )
 
 

cirq/_compat_test.py

@@ -87,7 +87,11 @@ def f(a, b):
         return a + b
 
     with cirq.testing.assert_deprecated(
-        'test_func was used', 'will be removed in cirq v1.2', 'Roll some dice.', deadline='v1.2'
+        '_compat_test.py:',
+        'test_func was used',
+        'will be removed in cirq v1.2',
+        'Roll some dice.',
+        deadline='v1.2',
     ):
         assert f(1, 2) == 3
 
@@ -101,7 +105,11 @@ def old_func(*args, **kwargs):
         return new_func(*args, **kwargs)
 
     with cirq.testing.assert_deprecated(
-        'old_func was used', 'will be removed in cirq v1.2', 'Roll some dice.', deadline='v1.2'
+        '_compat_test.py:',
+        'old_func was used',
+        'will be removed in cirq v1.2',
+        'Roll some dice.',
+        deadline='v1.2',
     ):
         assert old_func(1, 2) == 3
 
@@ -136,6 +144,7 @@ def f(new_count):
         assert f(new_count=1) == 1
 
     with cirq.testing.assert_deprecated(
+        '_compat_test.py:',
         'double_count parameter of test_func was used',
         'will be removed in cirq v1.2',
         'Double it yourself.',
@@ -194,6 +203,7 @@ def test_wrap_module():
 
     # Deprecation capability.
     with cirq.testing.assert_deprecated(
+        '_compat_test.py:',
         'foo was used but is deprecated.',
         'will be removed in cirq v0.6',
         'use bar instead',
@@ -232,6 +242,7 @@ class OldClass(NewClass):
     assert "OldClass docs" in OldClass.__doc__
 
     with cirq.testing.assert_deprecated(
+        '_compat_test.py:',
         'foo was used but is deprecated',
         'will be removed in cirq v1.2',
         'theFix',

cirq/protocols/json_serialization_test.py

@@ -518,7 +518,7 @@ def test_to_from_json_gzip():
 
 def _eval_repr_data_file(path: pathlib.Path, deprecation_deadline: Optional[str]):
     ctx_manager = (
-        cirq.testing.assert_deprecated(deadline=deprecation_deadline, allow_multiple_warnings=True)
+        cirq.testing.assert_deprecated(deadline=deprecation_deadline, count=None)
         if deprecation_deadline
         else contextlib.suppress()
     )
@@ -554,9 +554,7 @@ def assert_repr_and_json_test_data_agree(
     try:
         json_from_file = json_path.read_text()
         ctx_manager = (
-            cirq.testing.assert_deprecated(
-                deadline=deprecation_deadline, allow_multiple_warnings=True
-            )
+            cirq.testing.assert_deprecated(deadline=deprecation_deadline, count=None)
             if deprecation_deadline
             else contextlib.suppress()
         )

cirq/testing/deprecation.py

@@ -13,14 +13,29 @@
 # limitations under the License.
 import os
 from contextlib import contextmanager
+from typing import Optional
 
+from cirq._compat import deprecated_parameter
 from cirq.testing import assert_logs
 
 ALLOW_DEPRECATION_IN_TEST = 'ALLOW_DEPRECATION_IN_TEST'
 
 
 @contextmanager
-def assert_deprecated(*msgs: str, deadline: str, allow_multiple_warnings: bool = False):
+@deprecated_parameter(
+    deadline='v0.12',
+    fix='Use count instead.',
+    parameter_desc='allow_multiple_warnings',
+    match=lambda args, kwargs: 'allow_multiple_warnings' in kwargs,
+    rewrite=lambda args, kwargs: (
+        args,
+        dict(
+            ('count', None if v == True else 1) if k == 'allow_multiple_warnings' else (k, v)
+            for k, v in kwargs.items()
+        ),
+    ),
+)
+def assert_deprecated(*msgs: str, deadline: str, count: Optional[int] = 1):
     """Allows deprecated functions, classes, decorators in tests.
 
     It acts as a contextmanager that can be used in with statements:
@@ -31,13 +46,17 @@ def assert_deprecated(*msgs: str, deadline: str, allow_multiple_warnings: bool =
         msgs: messages that should match the warnings captured
         deadline: the expected deadline the feature will be deprecated by. Has to follow the format
             vX.Y (minor versions only)
-        allow_multiple_warnings: if True, multiple warnings are accepted. Typically this should not
-            be used, by default it's False.
+        count: if None count of messages is not asserted, otherwise the number of deprecation
+            messages have to equal count.
     """
 
     os.environ[ALLOW_DEPRECATION_IN_TEST] = 'True'
     try:
-        with assert_logs(*(msgs + (deadline,)), count=None if allow_multiple_warnings else 1):
+        with assert_logs(*(msgs + (deadline,)), count=count):
             yield True
     finally:
-        del os.environ[ALLOW_DEPRECATION_IN_TEST]
+        try:
+            del os.environ[ALLOW_DEPRECATION_IN_TEST]
+        except:
+            # this is only for nested deprecation checks
+            pass

cirq/testing/deprecation_test.py

@@ -36,6 +36,29 @@ def test_assert_deprecated_log_handling():
 
     # allowing for multiple deprecation warnings (in case of json serialization for multiple objects
     # for example)
-    with assert_deprecated(deadline="v1.2", allow_multiple_warnings=True):
+    with assert_deprecated(deadline="v1.2", count=2):
         warnings.warn("hello, this is deprecated in v1.2")
         warnings.warn("hello, this is deprecated in v1.2")
+
+    with assert_deprecated(deadline="v1.2", count=None):
+        warnings.warn("hello, this is deprecated in v1.2")
+        warnings.warn("hello, this is deprecated in v1.2")
+        warnings.warn("hello, this is deprecated in v1.2")
+
+
+def test_deprecated():
+    # allow_multiple_warnings is now deprecated...so this is a bit convoluted,
+    # a parameter of the deprecator is being deprecated
+
+    with assert_deprecated(deadline="v0.12", count=3):
+        with pytest.raises(AssertionError, match="Expected 1 log message but got 2."):
+            # pylint: disable=unexpected-keyword-arg
+            with assert_deprecated(deadline="v1.2", allow_multiple_warnings=False):
+                warnings.warn("hello, this is deprecated in v1.2")
+                warnings.warn("hello, this is deprecated in v1.2")
+
+    with assert_deprecated(deadline="v0.12", count=3):
+        # pylint: disable=unexpected-keyword-arg
+        with assert_deprecated(deadline="v1.2", allow_multiple_warnings=True):
+            warnings.warn("hello, this is deprecated in v1.2")
+            warnings.warn("hello, this is deprecated in v1.2")