Detect and support module aliasing via assignment. (#3435)

* Detect and support module aliasing via assignment.

In semantic analysis, if we find a simple assignment statement where the rvalue
is a module, make the lvalue a direct alias to that same module.

Fixes #1778.

* Handle chained assignment and iterable unpacking assignment.

* Remove test case that was only for exploration, not intended for inclusion.

* Add some more tests for module assignment.

* Also add tests for access/assignment of nonexistent module attribute.

* Break down code and add comments for clarity; add test for mismatch lengths.

* Naming improvements.

* Support tracking module assignment in non-global scope.

* Add more tests for unpacking mismatch cases.

* Keep rvals always on the right.

* Don't use a form of unpacking that is Py35+ only.

* It's the zip that is problematic, not just the unpacking.

* Add tests for module assignment in class and local scopes.

* Simplify to single method.

* Go back to annotating genericpath as Any in stdlib-sample.

* Respect explicit type annotation and don't propagate module reference.

* Backtrack to simple ModuleType var in case of incompatible module aliasing.

* Remove stray pdb comment.

* Also handle reassignment of an original (non-alias) module reference.

* Simplify: fail on assignment of different modules to same variable without explicit annotation.

* Style and working tweaks.

Author: carljm

mypy/semanal.py

@@ -1547,6 +1547,8 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.process_namedtuple_definition(s)
         self.process_typeddict_definition(s)
         self.process_enum_call(s)
+        if not s.type:
+            self.process_module_assignment(s.lvalues, s.rvalue, s)
 
         if (len(s.lvalues) == 1 and isinstance(s.lvalues[0], NameExpr) and
                 s.lvalues[0].name == '__all__' and s.lvalues[0].kind == GDEF and
@@ -2383,6 +2385,66 @@ def is_classvar(self, typ: Type) -> bool:
     def fail_invalid_classvar(self, context: Context) -> None:
         self.fail('ClassVar can only be used for assignments in class body', context)
 
+    def process_module_assignment(self, lvals: List[Expression], rval: Expression,
+                                  ctx: AssignmentStmt) -> None:
+        """Propagate module references across assignments.
+
+        Recursively handles the simple form of iterable unpacking; doesn't
+        handle advanced unpacking with *rest, dictionary unpacking, etc.
+
+        In an expression like x = y = z, z is the rval and lvals will be [x,
+        y].
+
+        """
+        if all(isinstance(v, (TupleExpr, ListExpr)) for v in lvals + [rval]):
+            # rval and all lvals are either list or tuple, so we are dealing
+            # with unpacking assignment like `x, y = a, b`. Mypy didn't
+            # understand our all(isinstance(...)), so cast them as
+            # Union[TupleExpr, ListExpr] so mypy knows it is safe to access
+            # their .items attribute.
+            seq_lvals = cast(List[Union[TupleExpr, ListExpr]], lvals)
+            seq_rval = cast(Union[TupleExpr, ListExpr], rval)
+            # given an assignment like:
+            #     (x, y) = (m, n) = (a, b)
+            # we now have:
+            #     seq_lvals = [(x, y), (m, n)]
+            #     seq_rval = (a, b)
+            # We now zip this into:
+            #     elementwise_assignments = [(a, x, m), (b, y, n)]
+            # where each elementwise assignment includes one element of rval and the
+            # corresponding element of each lval. Basically we unpack
+            #     (x, y) = (m, n) = (a, b)
+            # into elementwise assignments
+            #     x = m = a
+            #     y = n = b
+            # and then we recursively call this method for each of those assignments.
+            # If the rval and all lvals are not all of the same length, zip will just ignore
+            # extra elements, so no error will be raised here; mypy will later complain
+            # about the length mismatch in type-checking.
+            elementwise_assignments = zip(seq_rval.items, *[v.items for v in seq_lvals])
+            for rv, *lvs in elementwise_assignments:
+                self.process_module_assignment(lvs, rv, ctx)
+        elif isinstance(rval, NameExpr):
+            rnode = self.lookup(rval.name, ctx)
+            if rnode and rnode.kind == MODULE_REF:
+                for lval in lvals:
+                    if not isinstance(lval, NameExpr):
+                        continue
+                    # respect explicitly annotated type
+                    if (isinstance(lval.node, Var) and lval.node.type is not None):
+                        continue
+                    lnode = self.lookup(lval.name, ctx)
+                    if lnode:
+                        if lnode.kind == MODULE_REF and lnode.node is not rnode.node:
+                            self.fail(
+                                "Cannot assign multiple modules to name '{}' "
+                                "without explicit 'types.ModuleType' annotation".format(lval.name),
+                                ctx)
+                        # never create module alias except on initial var definition
+                        elif lval.is_def:
+                            lnode.kind = MODULE_REF
+                            lnode.node = rnode.node
+
     def process_enum_call(self, s: AssignmentStmt) -> None:
         """Check if s defines an Enum; if yes, store the definition in symbol table."""
         if len(s.lvalues) != 1 or not isinstance(s.lvalues[0], NameExpr):

test-data/stdlib-samples/3.2/test/test_genericpath.py

@@ -23,7 +23,7 @@ def safe_rmdir(dirname: str) -> None:
 
 class GenericTest(unittest.TestCase):
     # The path module to be tested
-    pathmodule = genericpath # type: Any
+    pathmodule = genericpath  # type: Any
     common_attributes = ['commonprefix', 'getsize', 'getatime', 'getctime',
                          'getmtime', 'exists', 'isdir', 'isfile']
     attributes = []  # type: List[str]

test-data/unit/check-modules.test

@@ -1439,3 +1439,204 @@ class C:
 a = 'foo'
 
 [builtins fixtures/module.pyi]
+
+[case testModuleAlias]
+import m
+m2 = m
+reveal_type(m2.a)  # E: Revealed type is 'builtins.str'
+m2.b  # E: Module has no attribute "b"
+m2.c = 'bar'  # E: Module has no attribute "c"
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testClassModuleAlias]
+import m
+
+class C:
+    x = m
+    def foo(self) -> None:
+        reveal_type(self.x.a)  # E: Revealed type is 'builtins.str'
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testLocalModuleAlias]
+import m
+
+def foo() -> None:
+    x = m
+    reveal_type(x.a)  # E: Revealed type is 'builtins.str'
+
+class C:
+    def foo(self) -> None:
+        x = m
+        reveal_type(x.a)  # E: Revealed type is 'builtins.str'
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testChainedModuleAlias]
+import m
+m3 = m2 = m
+m4 = m3
+m5 = m4
+reveal_type(m2.a)  # E: Revealed type is 'builtins.str'
+reveal_type(m3.a)  # E: Revealed type is 'builtins.str'
+reveal_type(m4.a)  # E: Revealed type is 'builtins.str'
+reveal_type(m5.a)  # E: Revealed type is 'builtins.str'
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testMultiModuleAlias]
+import m, n
+m2, n2, (m3, n3) = m, n, [m, n]
+reveal_type(m2.a)  # E: Revealed type is 'builtins.str'
+reveal_type(n2.b)  # E: Revealed type is 'builtins.str'
+reveal_type(m3.a)  # E: Revealed type is 'builtins.str'
+reveal_type(n3.b)  # E: Revealed type is 'builtins.str'
+
+x, y = m  # E: 'types.ModuleType' object is not iterable
+x, y, z = m, n  # E: Need more than 2 values to unpack (3 expected)
+x, y = m, m, m  # E: Too many values to unpack (2 expected, 3 provided)
+x, (y, z) = m, n  # E: 'types.ModuleType' object is not iterable
+x, (y, z) = m, (n, n, n)  # E: Too many values to unpack (2 expected, 3 provided)
+
+[file m.py]
+a = 'foo'
+
+[file n.py]
+b = 'bar'
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasWithExplicitAnnotation]
+from typing import Any
+import types
+import m
+mod_mod: types.ModuleType = m
+mod_mod2: types.ModuleType
+mod_mod2 = m
+mod_mod3 = m  # type: types.ModuleType
+mod_any: Any = m
+mod_int: int = m  # E: Incompatible types in assignment (expression has type Module, variable has type "int")
+
+reveal_type(mod_mod)  # E: Revealed type is 'types.ModuleType'
+mod_mod.a  # E: Module has no attribute "a"
+reveal_type(mod_mod2)  # E: Revealed type is 'types.ModuleType'
+mod_mod2.a  # E: Module has no attribute "a"
+reveal_type(mod_mod3)  # E: Revealed type is 'types.ModuleType'
+mod_mod3.a  # E: Module has no attribute "a"
+reveal_type(mod_any)  # E: Revealed type is 'Any'
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasPassedToFunction]
+import types
+import m
+
+def takes_module(x: types.ModuleType):
+    reveal_type(x.__file__)  # E: Revealed type is 'builtins.str'
+
+n = m
+takes_module(m)
+takes_module(n)
+
+[file m.py]
+a = 'foo'
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasRepeated]
+import m, n
+
+if bool():
+    x = m
+else:
+    x = 3  # E: Incompatible types in assignment (expression has type "int", variable has type Module)
+
+if bool():
+    y = 3
+else:
+    y = m  # E: Incompatible types in assignment (expression has type Module, variable has type "int")
+
+if bool():
+    z = m
+else:
+    z = n  # E: Cannot assign multiple modules to name 'z' without explicit 'types.ModuleType' annotation
+
+[file m.py]
+a = 'foo'
+
+[file n.py]
+a = 3
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasRepeatedWithAnnotation]
+import types
+import m, n
+
+x: types.ModuleType
+if bool():
+    x = m
+else:
+    x = n
+
+x.a  # E: Module has no attribute "a"
+reveal_type(x.__file__)  # E: Revealed type is 'builtins.str'
+
+[file m.py]
+a = 'foo'
+
+[file n.py]
+a = 3
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasRepeatedComplex]
+import m, n, o
+
+x = m
+x = n  # E: Cannot assign multiple modules to name 'x' without explicit 'types.ModuleType' annotation
+x = o  # E: Cannot assign multiple modules to name 'x' without explicit 'types.ModuleType' annotation
+
+y = o
+y, z = m, n  # E: Cannot assign multiple modules to name 'y' without explicit 'types.ModuleType' annotation
+
+xx = m
+xx = m
+reveal_type(xx.a)  # E: Revealed type is 'builtins.str'
+
+[file m.py]
+a = 'foo'
+
+[file n.py]
+a = 3
+
+[file o.py]
+a = 'bar'
+
+[builtins fixtures/module.pyi]
+
+[case testModuleAliasToOtherModule]
+import m, n
+m = n  # E: Cannot assign multiple modules to name 'm' without explicit 'types.ModuleType' annotation
+
+[file m.py]
+
+[file n.py]
+
+[builtins fixtures/module.pyi]

test-data/unit/lib-stub/types.pyi

@@ -6,4 +6,5 @@ def coroutine(func: _T) -> _T: pass
 
 class bool: ...
 
-class ModuleType: ...
+class ModuleType:
+    __file__ = ... # type: str