Add tests for module wrapping in _compat

cirq/_compat.py

@@ -185,17 +185,18 @@ def decorated_func(*args, **kwargs) -> Any:
 
 def wrap_module(module: ModuleType,
                 deprecated_attributes: Dict[str, Tuple[str, str]]):
-    """Wrap a module with deprecated attributes.
+    """Wrap a module with deprecated attributes that give warnings.
 
     Args:
         module: The module to wrap.
-        deprecated_attributes: A dictionary from attribute name to pair of
+        deprecated_attributes: A dictionary from attribute name to a tuple of
             strings, where the first string gives the version that the attribute
             will be removed in, and the second string describes what the user
             should do instead of accessing this deprecated attribute.
 
     Returns:
-        Wrapped module with deprecated attributes.
+        Wrapped module with deprecated attributes. Use of these attributes
+        will cause a warning for these deprecated attributes.
     """
 
     class Wrapped(ModuleType):
@@ -213,4 +214,4 @@ def __getattr__(self, name):
                     stacklevel=2)
             return getattr(module, name)
 
-    return Wrapped(module.__name__)
+    return Wrapped(module.__name__, module.__doc__)

cirq/_compat_test.py

@@ -13,18 +13,15 @@
 # limitations under the License.
 
 import logging
+import types
 from typing import ContextManager, List
 
 import numpy as np
 import pandas as pd
 import sympy
 
-from cirq._compat import (
-    proper_repr,
-    deprecated,
-    deprecated_parameter,
-    proper_eq,
-)
+from cirq._compat import (proper_repr, deprecated, deprecated_parameter,
+                          proper_eq, wrap_module)
 
 
 def test_proper_repr():
@@ -122,6 +119,37 @@ def f(new_count):
     assert 'Double it yourself.' in log[0].getMessage()
 
 
+def test_wrap_module():
+    my_module = types.ModuleType('my_module', 'my doc string')
+    my_module.foo = 'foo'
+    my_module.bar = 'bar'
+    assert 'foo' in my_module.__dict__
+    assert 'bar' in my_module.__dict__
+    assert 'zoo' not in my_module.__dict__
+
+    wrapped = wrap_module(my_module, {'foo': ('0.6.0', 'use bar instead')})
+    # Dunder methods
+    assert wrapped.__doc__ == 'my doc string'
+    assert wrapped.__name__ == 'my_module'
+    # Test dict is correct.
+    assert 'foo' in wrapped.__dict__
+    assert 'bar' in wrapped.__dict__
+    assert 'zoo' not in wrapped.__dict__
+
+    # Deprecation capability.
+    with capture_logging() as log:
+        _ = wrapped.foo
+    assert len(log) == 1
+    msg = log[0].getMessage()
+    assert 'foo was used but is deprecated.' in msg
+    assert 'will be removed in cirq 0.6.0' in msg
+    assert 'use bar instead' in msg
+
+    with capture_logging() as log:
+        _ = wrapped.bar
+    assert len(log) == 0
+
+
 def capture_logging() -> ContextManager[List[logging.LogRecord]]:
     records = []