Detect and support module aliasing via assignment.

mypy/semanal.py

@@ -1520,6 +1520,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.process_namedtuple_definition(s)
         self.process_typeddict_definition(s)
         self.process_enum_call(s)
+        self.process_module_assignment(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
@@ -2356,6 +2357,64 @@ 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, s: AssignmentStmt) -> None:
+        """Check if s assigns a module an alias name; if so, update symbol table."""
+        self._process_module_assignment(s.lvalues, s.rvalue, s)
+
+    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
+                    lnode = self.lookup(lval.name, ctx)
+                    if lnode:
+                        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
     common_attributes = ['commonprefix', 'getsize', 'getatime', 'getctime',
                          'getmtime', 'exists', 'isdir', 'isfile']
     attributes = []  # type: List[str]
@@ -143,7 +143,7 @@ def test_exists(self) -> None:
             f.close()
             self.assertIs(self.pathmodule.exists(support.TESTFN), True)
             if not self.pathmodule == genericpath:
-                self.assertIs(self.pathmodule.lexists(support.TESTFN),
+                self.assertIs(self.pathmodule.lexists(support.TESTFN),  # type: ignore
                               True)
         finally:
             if not f.closed:

test-data/unit/check-modules.test

@@ -1439,3 +1439,82 @@ class C:
 a = 'foo'
 
 [builtins fixtures/module.pyi]
+
+[case testModuleAssignment]
+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 testClassModuleAssignment]
+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 testLocalModuleAssignment]
+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 testChainedModuleAssignment]
+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 testMultiModuleAssignment]
+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]