Merge pull request python#466 from spkersten/nestup248

Assignment to and from nested tuples (or lists)

mypy/checkexpr.py

@@ -1188,7 +1188,7 @@ def check_generator_or_comprehension(self, gen: GeneratorExpr,
         for index, sequence, conditions in zip(gen.indices, gen.sequences,
                                                gen.condlists):
             sequence_type = self.chk.analyse_iterable_item_type(sequence)
-            self.chk.analyse_index_variables(index, False, sequence_type, gen)
+            self.chk.analyse_index_variables(index, sequence_type, gen)
             for condition in conditions:
                 self.accept(condition)
         self.chk.binder.pop_frame()

mypy/messages.py

@@ -454,6 +454,18 @@ def invalid_cast(self, target_type: Type, source_type: Type,
             self.fail('Cannot cast from {} to {}'.format(
                 self.format(source_type), self.format(target_type)), context)
 
+    def wrong_number_values_to_unpack(self, provided: int, expected: int, context: Context) -> None:
+        if provided < expected:
+            if provided == 1:
+                self.fail('Need more than 1 value to unpack ({} expected)'.format(expected), context)
+            else:
+                self.fail('Need more than {} values to unpack ({} expected)'.format(provided, expected), context)
+        elif provided > expected:
+            self.fail('Too many values to unpack ({} expected, {} provided)'.format(expected, provided), context)
+
+    def type_not_iterable(self, type: Type, context: Context) -> None:
+        self.fail('\'{}\' object is not iterable'.format(type), context)
+
     def incompatible_operator_assignment(self, op: str,
                                          context: Context) -> None:
         self.fail('Result type of {} incompatible in assignment'.format(op),

mypy/nodes.py

@@ -569,32 +569,21 @@ def accept(self, visitor: NodeVisitor[T]) -> T:
 class ForStmt(Node):
     # Index variables
     index = Undefined(List['NameExpr'])
-    # Index variable types (each may be None)
-    types = Undefined(List['mypy.types.Type'])
     # Expression to iterate
     expr = Undefined(Node)
     body = Undefined(Block)
     else_body = Undefined(Block)
 
     def __init__(self, index: List['NameExpr'], expr: Node, body: Block,
-                 else_body: Block,
-                 types: List['mypy.types.Type'] = None) -> None:
+                 else_body: Block) -> None:
         self.index = index
         self.expr = expr
         self.body = body
         self.else_body = else_body
-        self.types = types
 
     def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit_for_stmt(self)
 
-    def is_annotated(self) -> bool:
-        ann = False
-        for t in self.types:
-            if t is not None:
-                ann = True
-        return ann
-
 
 class ReturnStmt(Node):
     expr = Undefined(Node)   # Expression or None

mypy/output.py

@@ -263,7 +263,6 @@ def visit_for_stmt(self, o):
         r = o.repr
         self.token(r.for_tok)
         for i in range(len(o.index)):
-            self.type(o.types[i])
             self.node(o.index[i])
             self.token(r.commas[i])
         self.token(r.in_tok)

mypy/parse.py

@@ -864,7 +864,7 @@ def parse_for_stmt(self) -> ForStmt:
             else_body = None
             else_tok = none
 
-        node = ForStmt(index, expr, body, else_body, types)
+        node = ForStmt(index, expr, body, else_body)
         self.set_repr(node, noderepr.ForStmtRepr(for_tok, commas, in_tok,
                                                  else_tok))
         return node

mypy/parsetype.py

@@ -51,6 +51,8 @@ def index(self) -> int:
     def parse_type(self) -> Type:
         """Parse a type."""
         t = self.current_token()
+        if t.string == '(':
+            return self.parse_parens()
         if isinstance(t, Name):
             return self.parse_named_type()
         elif t.string == '[':
@@ -68,20 +70,31 @@ def parse_type(self) -> Type:
         else:
             self.parse_error()
 
+    def parse_parens(self) -> Type:
+        self.expect('(')
+        types = self.parse_types()
+        self.expect(')')
+        return types
+
     def parse_types(self) -> Type:
-        parens = False
-        if self.current_token_str() == '(':
-            self.skip()
-            parens = True
+        """ Parse either a single type or a comma separated
+        list of types as a tuple type. In the latter case, a
+        trailing comma is needed when the list contains only
+        a single type (and optional otherwise).
+
+        int   ->   int
+        int,  ->   TupleType[int]
+        int, int, int  ->  TupleType[int, int, int]
+        """
         type = self.parse_type()
         if self.current_token_str() == ',':
             items = [type]
             while self.current_token_str() == ',':
                 self.skip()
+                if self.current_token_str() == ')':
+                    break
                 items.append(self.parse_type())
             type = TupleType(items, None)
-        if parens:
-            self.expect(')')
         return type
 
     def parse_type_list(self) -> TypeList:

mypy/semanal.py

@@ -735,6 +735,7 @@ def analyse_lvalue(self, lval: Node, nested: bool = False,
         Only if add_global is True, add name to globals table. If nested
         is true, the lvalue is within a tuple or list lvalue expression.
         """
+
         if isinstance(lval, NameExpr):
             nested_global = (not self.is_func_scope() and
                              self.block_depth[-1] > 0 and
@@ -795,7 +796,7 @@ def analyse_lvalue(self, lval: Node, nested: bool = False,
         elif isinstance(lval, ParenExpr):
             self.analyse_lvalue(lval.expr, nested, add_global, explicit_type)
         elif (isinstance(lval, TupleExpr) or
-              isinstance(lval, ListExpr)) and not nested:
+              isinstance(lval, ListExpr)):
             items = (Any(lval)).items
             for i in items:
                 self.analyse_lvalue(i, nested=True, add_global=add_global,
@@ -1027,19 +1028,6 @@ def visit_for_stmt(self, s: ForStmt) -> None:
         for n in s.index:
             self.analyse_lvalue(n)
 
-        # Analyze index variable types.
-        for i in range(len(s.types)):
-            t = s.types[i]
-            if t:
-                s.types[i] = self.anal_type(t)
-                v = cast(Var, s.index[i].node)
-                # TODO check if redefinition
-                v.type = s.types[i]
-
-        # Report error if only some of the loop variables have annotations.
-        if s.types != [None] * len(s.types) and None in s.types:
-            self.fail('Cannot mix unannotated and annotated loop variables', s)
-
         self.loop_depth += 1
         self.visit_block(s.body)
         self.loop_depth -= 1

mypy/strconv.py

@@ -197,8 +197,6 @@ def visit_while_stmt(self, o):
 
     def visit_for_stmt(self, o):
         a = [o.index]
-        if o.types != [None] * len(o.types):
-            a += o.types
         a.extend([o.expr, o.body])
         if o.else_body:
             a.append(('Else', o.else_body.body))

mypy/test/data/check-inference.test

@@ -194,6 +194,78 @@ class B: pass
 [out]
 main: In function "f":
 
+[case testInferringLvarTypesInTupleAssignment]
+from typing import Undefined, Tuple
+def f() -> None:
+    t = Undefined # type: Tuple[A, B]
+    a, b = t
+    a = b   # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    a = B() # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    b = A() # E: Incompatible types in assignment (expression has type "A", variable has type "B")
+
+    a = A()
+    b = B()
+
+class A: pass
+class B: pass
+[out]
+main: In function "f":
+
+[case testInferringLvarTypesInNestedTupleAssignment1]
+from typing import Undefined, Tuple
+def f() -> None:
+    t = Undefined # type: Tuple[A, B]
+    a1, (a, b) = A(), t
+    a = b   # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    a = B() # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    b = A() # E: Incompatible types in assignment (expression has type "A", variable has type "B")
+
+    a = A()
+    b = B()
+
+class A: pass
+class B: pass
+[out]
+main: In function "f":
+
+[case testInferringLvarTypesInNestedTupleAssignment2]
+import typing
+def f() -> None:
+    a, (b, c) = A(), (B(), C())
+    a = b   # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    a = B() # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    b = A() # E: Incompatible types in assignment (expression has type "A", variable has type "B")
+    c = A() # E: Incompatible types in assignment (expression has type "A", variable has type "C")
+
+    a = A()
+    b = B()
+    c = C()
+
+class A: pass
+class B: pass
+class C: pass
+[out]
+main: In function "f":
+
+[case testInferringLvarTypesInNestedListAssignment]
+import typing
+def f() -> None:
+    a, (b, c) = A(), [B(), C()]
+    a = b   # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    a = B() # E: Incompatible types in assignment (expression has type "B", variable has type "A")
+    b = A() # E: Incompatible types in assignment (expression has type "A", variable has type "B")
+    c = A() # E: Incompatible types in assignment (expression has type "A", variable has type "C")
+
+    a = A()
+    b = B()
+    c = C()
+
+class A: pass
+class B: pass
+class C: pass
+[out]
+main: In function "f":
+
 [case testInferringLvarTypesInMultiDefWithNoneTypes]
 import typing
 def f() -> None:
@@ -204,6 +276,15 @@ class A: pass
 [out]
 main: In function "f":
 
+[case testInferringLvarTypesInNestedTupleAssignmentWithNoneTypes]
+import typing
+def f() -> None:
+    a1, (a2, b) = A(), (A(), None) # E: Need type annotation for variable
+
+class A: pass
+[out]
+main: In function "f":
+
 [case testInferringLvarTypesInMultiDefWithInvalidTuple]
 from typing import Undefined, Tuple
 t = Undefined # type: Tuple[object, object, object]
@@ -215,14 +296,24 @@ def f() -> None:
 [builtins fixtures/tuple.py]
 [out]
 main: In function "f":
-main, line 5: Too many values to assign
-main, line 6: Need 4 values to assign
+main, line 5: Too many values to unpack (2 expected, 3 provided)
+main, line 6: Need more than 3 values to unpack (4 expected)
 
 [case testInvalidRvalueTypeInInferredMultipleLvarDefinition]
 import typing
 def f() -> None:
-    a, b = f   # E: Incompatible types in assignment
-    c, d = A() # E: Incompatible types in assignment
+    a, b = f   # E: 'def ()' object is not iterable
+    c, d = A() # E: '__main__.A' object is not iterable
+class A: pass
+[builtins fixtures/for.py]
+[out]
+main: In function "f":
+
+[case testInvalidRvalueTypeInInferredNestedTupleAssignment]
+import typing
+def f() -> None:
+    a1, (a2, b) = A(), f   # E: 'def ()' object is not iterable
+    a3, (c, d) = A(), A() # E: '__main__.A' object is not iterable
 class A: pass
 [builtins fixtures/for.py]
 [out]
@@ -255,6 +346,33 @@ def f() -> None:
 [out]
 main: In function "f":
 
+[case testInferringNestedTupleAssignmentWithListRvalue]
+from typing import List
+
+class C: pass
+class D: pass
+
+def f() -> None:
+    c1, (a, b) = C(), List[C]()
+    c2, (c, d, e) = C(), List[D]()
+    a = D() # E: Incompatible types in assignment (expression has type "D", variable has type "C")
+    b = D() # E: Incompatible types in assignment (expression has type "D", variable has type "C")
+    c = C() # E: Incompatible types in assignment (expression has type "C", variable has type "D")
+    b = c   # E: Incompatible types in assignment (expression has type "D", variable has type "C")
+
+    a = C()
+    b = C()
+    c = D()
+    d = D()
+    e = D()
+
+    a = b
+    c = d
+    d = e
+[builtins fixtures/for.py]
+[out]
+main: In function "f":
+
 [case testInferringMultipleLvarDefinitionWithImplicitDynamicRvalue]
 import typing
 def f() -> None:
@@ -639,7 +757,7 @@ class B: pass
 [out]
 main, line 4: Incompatible types in assignment (expression has type "A", variable has type "B")
 main, line 5: Incompatible types in assignment (expression has type "B", variable has type "A")
-main, line 8: Need 3 values to assign
+main, line 8: Need more than 2 values to unpack (3 expected)
 main, line 10: Need type annotation for variable
 
 [case testInferenceOfFor3]
@@ -711,7 +829,7 @@ class A: pass
 [case testMultipleAssignmentWithPartialDefinition2]
 from typing import Undefined
 a = Undefined # type: A
-a, x = [a]
+a, x = [a, a]
 x = a
 a = x
 x = object() # E: Incompatible types in assignment (expression has type "object", variable has type "A")