Handle chained assignment and iterable unpacking assignment.

Author: carljm

mypy/semanal.py

@@ -2357,22 +2357,39 @@ 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 yes, update symbol table."""
-        # TODO support more complex forms of module alias assignment
-        # (e.g. `x, y = (mod1, mod2)`) and aliases not in global scope
-        if (
-                len(s.lvalues) != 1
-                or not isinstance(s.lvalues[0], NameExpr)
-                or not isinstance(s.rvalue, NameExpr)
-                or not self.is_module_scope()
-        ):
+        """Check if s assigns a module an alias name; if so, update symbol table."""
+        # TODO support module alias assignment not in global scope
+        if not self.is_module_scope():
             return
-        rnode = self.lookup(s.rvalue.name, s)
-        if rnode and rnode.kind == MODULE_REF:
-            lnode = self.lookup(s.lvalues[0].name, s)
-            if lnode:
-                lnode.kind = MODULE_REF
-                lnode.node = rnode.node
+        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.
+
+        """
+        if all(isinstance(v, (TupleExpr, ListExpr)) for v in lvals + [rval]):
+            litemlists = [v.items for v in cast(List[Union[TupleExpr, ListExpr]], lvals)]
+            ritems = cast(Union[TupleExpr, ListExpr], rval).items
+            for rv, *lvs in zip(ritems, *litemlists):
+                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."""

test-data/unit/check-modules.test

@@ -1426,3 +1426,37 @@ reveal_type(m2.a)  # E: Revealed type is 'builtins.str'
 a = 'foo'
 
 [builtins fixtures/module.pyi]
+
+[case testChainedModuleAssignment]
+import m
+m3 = m2 = m
+reveal_type(m2.a)  # E: Revealed type is 'builtins.str'
+reveal_type(m3.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'
+
+[file m.py]
+a = 'foo'
+
+[file n.py]
+b = 'bar'
+
+[builtins fixtures/module.pyi]
+
+[case testMultiAssignment]
+m2, n2 = 'foo', 'bar'
+reveal_type(m2)  # E: Revealed type is 'builtins.str'
+reveal_type(n2)  # E: Revealed type is 'builtins.str'
+
+[builtins fixtures/module.pyi]