Fix submodule attributes and spec after deprecation (#4920)

* add `__spec__` to the wrapper module produced by deprecated_submodule and
  deprecate_attributes and make it proxy the spec of the original module.
  This is required for later submodule imports.

* when deprecated_submodule and deprecate_attributes replace items in
  sys.modules update submodule attributes in their parents to be the same.

* make deprecate_attributes more convenient by replacing the module
  argument with the absolute module name.  Let it then handle sys.modules and
  attributes updates.

cirq-core/cirq/_compat.py

@@ -328,11 +328,11 @@ def decorated_func(*args, **kwargs) -> Any:
     return decorator
 
 
-def deprecate_attributes(module: ModuleType, deprecated_attributes: Dict[str, Tuple[str, str]]):
-    """Wrap a module with deprecated attributes that give warnings.
+def deprecate_attributes(module_name: str, deprecated_attributes: Dict[str, Tuple[str, str]]):
+    """Replace module with a wrapper that gives warnings for deprecated attributes.
 
     Args:
-        module: The module to wrap.
+        module_name: Absolute name of the module that deprecates attributes.
         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
@@ -346,10 +346,15 @@ def deprecate_attributes(module: ModuleType, deprecated_attributes: Dict[str, Tu
     for (deadline, _) in deprecated_attributes.values():
         _validate_deadline(deadline)
 
+    module = sys.modules[module_name]
+
     class Wrapped(ModuleType):
 
         __dict__ = module.__dict__
 
+        # Workaround for: https://github.com/python/mypy/issues/8083
+        __spec__ = _make_proxy_spec_property(module)  # type: ignore
+
         def __getattr__(self, name):
             if name in deprecated_attributes:
                 deadline, fix = deprecated_attributes[name]
@@ -360,7 +365,13 @@ def __getattr__(self, name):
                 )
             return getattr(module, name)
 
-    return Wrapped(module.__name__, module.__doc__)
+    wrapped_module = Wrapped(module_name, module.__doc__)
+    if '.' in module_name:
+        parent_name, module_tail = module_name.rsplit('.', 1)
+        setattr(sys.modules[parent_name], module_tail, wrapped_module)
+    sys.modules[module_name] = wrapped_module
+
+    return wrapped_module
 
 
 class DeprecatedModuleLoader(importlib.abc.Loader):
@@ -642,14 +653,32 @@ def _setup_deprecated_submodule_attribute(
     class Wrapped(ModuleType):
         __dict__ = parent_module.__dict__
 
+        # Workaround for: https://github.com/python/mypy/issues/8083
+        __spec__ = _make_proxy_spec_property(parent_module)  # type: ignore
+
         def __getattr__(self, name):
             if name == old_child:
                 _deduped_module_warn_or_error(
                     f"{old_parent}.{old_child}", new_module_name, deadline
                 )
             return getattr(parent_module, name)
 
-    sys.modules[old_parent] = Wrapped(parent_module.__name__, parent_module.__doc__)
+    wrapped_parent_module = Wrapped(parent_module.__name__, parent_module.__doc__)
+    if '.' in old_parent:
+        grandpa_name, parent_tail = old_parent.rsplit('.', 1)
+        grandpa_module = sys.modules[grandpa_name]
+        setattr(grandpa_module, parent_tail, wrapped_parent_module)
+    sys.modules[old_parent] = wrapped_parent_module
+
+
+def _make_proxy_spec_property(source_module: ModuleType) -> property:
+    def fget(self):
+        return source_module.__spec__
+
+    def fset(self, value):
+        source_module.__spec__ = value
+
+    return property(fget, fset)
 
 
 @contextlib.contextmanager

cirq-core/cirq/_compat_test.py

@@ -253,17 +253,25 @@ def test_wrap_module():
     my_module = types.ModuleType('my_module', 'my doc string')
     my_module.foo = 'foo'
     my_module.bar = 'bar'
+    my_module.__spec__ = ModuleSpec('my_module', loader=None)
     assert 'foo' in my_module.__dict__
     assert 'bar' in my_module.__dict__
     assert 'zoo' not in my_module.__dict__
 
     with pytest.raises(AssertionError, match='deadline should match vX.Y'):
         deprecate_attributes(my_module, {'foo': ('invalid', 'use bar instead')})
 
-    wrapped = deprecate_attributes(my_module, {'foo': ('v0.6', 'use bar instead')})
+    # temporarily update sys.modules so deprecate_attributes can find my_module
+    sys.modules['my_module'] = my_module
+    wrapped = deprecate_attributes('my_module', {'foo': ('v0.6', 'use bar instead')})
+    assert wrapped is sys.modules.pop('my_module')
     # Dunder methods
     assert wrapped.__doc__ == 'my doc string'
     assert wrapped.__name__ == 'my_module'
+    assert wrapped.__spec__ is my_module.__spec__
+    # Verify __spec__ setter in the wrapped module
+    wrapped.__spec__ = ModuleSpec('my_module', loader=None)
+    assert my_module.__spec__ is wrapped.__spec__
     # Test dict is correct.
     assert 'foo' in wrapped.__dict__
     assert 'bar' in wrapped.__dict__
@@ -288,6 +296,24 @@ def test_wrap_module():
         _ = wrapped.bar
 
 
+def test_deprecate_attributes_assert_attributes_in_sys_modules():
+    subprocess_context(_test_deprecate_attributes_assert_attributes_in_sys_modules)()
+
+
+def _test_deprecate_attributes_assert_attributes_in_sys_modules():
+    """Ensure submodule attributes are consistent with sys.modules items."""
+    import cirq.testing._compat_test_data.module_a as module_a0
+
+    module_a1 = deprecate_attributes(
+        'cirq.testing._compat_test_data.module_a',
+        {'MODULE_A_ATTRIBUTE': ('v0.6', 'use plain string instead')},
+    )
+
+    assert module_a1 is not module_a0
+    assert module_a1 is cirq.testing._compat_test_data.module_a
+    assert module_a1 is sys.modules['cirq.testing._compat_test_data.module_a']
+
+
 def test_deprecated_class():
     class NewClass:
         def __init__(self, a):
@@ -429,6 +455,21 @@ def _import_multiple_deprecated():
     assert module_c.MODULE_C_ATTRIBUTE == 'module_c'
 
 
+def _deprecate_grandchild_assert_attributes_in_sys_modules():
+    """Ensure submodule attributes are identical to sys.modules values."""
+    import cirq.testing._compat_test_data.module_a.fake_ab  # type: ignore
+
+    assert (
+        cirq.testing._compat_test_data.module_a.fake_ab
+        is sys.modules['cirq.testing._compat_test_data.module_a.fake_ab']
+    )
+    assert (
+        cirq.testing._compat_test_data.module_a
+        is sys.modules['cirq.testing._compat_test_data.module_a']
+    )
+    assert cirq.testing._compat_test_data is sys.modules['cirq.testing._compat_test_data']
+
+
 def _new_module_in_different_parent():
     from cirq.testing._compat_test_data.fake_ops import raw_types  # type: ignore
 
@@ -520,6 +561,12 @@ def _type_repr_in_deprecated_module():
     f'Use {old_parent}.module_a.module_b instead',
 ] + _deprecation_origin
 
+# see cirq_compat_test_data/__init__.py for the setup code
+_fake_ab_deprecation_msg = [
+    f'{old_parent}.module_a.fake_ab was used but is deprecated',
+    f'Use {old_parent}.module_a.module_b instead',
+] + _deprecation_origin
+
 # see cirq_compat_test_data/__init__.py for the setup code
 _fake_ops_deprecation_msg = [
     f'{old_parent}.fake_ops was used but is deprecated',
@@ -586,6 +633,7 @@ def isolated_func(*args, **kwargs):
         (_import_deprecated_first_new_second, [_fake_a_deprecation_msg]),
         (_import_new_first_deprecated_second, [_fake_a_deprecation_msg]),
         (_import_multiple_deprecated, [_fake_a_deprecation_msg, _fake_b_deprecation_msg]),
+        (_deprecate_grandchild_assert_attributes_in_sys_modules, [_fake_ab_deprecation_msg]),
         (_new_module_in_different_parent, [_fake_ops_deprecation_msg]),
         # ignore the frame requirement - as we are using find_spec from importlib, it
         # is detected as an 'internal' frame by warnings
@@ -691,6 +739,16 @@ def _test_metadata_distributions_after_deprecated_submodule():
     assert all(isinstance(d.name, str) for d in distlist)
 
 
+def test_parent_spec_after_deprecated_submodule():
+    subprocess_context(_test_parent_spec_after_deprecated_submodule)()
+
+
+def _test_parent_spec_after_deprecated_submodule():
+    import cirq.testing._compat_test_data
+
+    assert cirq.testing._compat_test_data.__spec__.name == 'cirq.testing._compat_test_data'
+
+
 def test_type_repr_in_new_module():
     # to cater for metadata path finders
     # https://docs.python.org/3/library/importlib.metadata.html#extending-the-search-algorithm

cirq-core/cirq/testing/_compat_test_data/__init__.py

@@ -19,6 +19,16 @@
     create_attribute=True,
 )
 
+# simulates a rename of a grandchild module
+# module_a.fake_ab -> module_a.module_b
+_compat.deprecated_submodule(
+    new_module_name=f"{__name__}.module_a.module_b",
+    old_parent=f"{__name__}.module_a",
+    old_child="fake_ab",
+    deadline="v0.20",
+    create_attribute=True,
+)
+
 # simulates a rename of a child module with same prefix
 # this prefix will collide with multiple "old" and new prefixes here
 # module_ -> module_a

cirq-google/cirq_google/__init__.py

@@ -153,12 +153,10 @@
 
 _register_resolver(_class_resolver_dictionary)
 
-__spec_copy__ = sys.modules[__name__].__spec__
-sys.modules[__name__] = _compat.deprecate_attributes(
-    sys.modules[__name__],
+_compat.deprecate_attributes(
+    __name__,
     {
         'Bristlecone': ('v0.15', 'Bristlecone will no longer be supported.'),
         'Foxtail': ('v0.15', 'Foxtail will no longer be supported.'),
     },
 )
-sys.modules[__name__].__spec__ = __spec_copy__