diff --git a/mypy/checker.py b/mypy/checker.py index b01be788aaa4..2558440168e3 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -2482,8 +2482,47 @@ def check_assignment_to_multiple_lvalues(self, lvalues: List[Lvalue], rvalue: Ex # using the type of rhs, because this allowed more fine grained # control in cases like: a, b = [int, str] where rhs would get # type List[object] - - rvalues = rvalue.items + rvalues = [] # type: List[Expression] + iterable_type = None # type: Optional[Type] + last_idx = None # type: Optional[int] + for idx_rval, rval in enumerate(rvalue.items): + if isinstance(rval, StarExpr): + typs = get_proper_type(self.expr_checker.visit_star_expr(rval).type) + if isinstance(typs, TupleType): + rvalues.extend([TempNode(typ) for typ in typs.items]) + elif self.type_is_iterable(typs) and isinstance(typs, Instance): + if (iterable_type is not None + and iterable_type != self.iterable_item_type(typs)): + self.fail("Contiguous iterable with same type expected", context) + else: + if last_idx is None or last_idx + 1 == idx_rval: + rvalues.append(rval) + last_idx = idx_rval + iterable_type = self.iterable_item_type(typs) + else: + self.fail("Contiguous iterable with same type expected", context) + else: + self.fail("Invalid type '{}' for *expr (iterable expected)".format(typs), + context) + else: + rvalues.append(rval) + iterable_start = None # type: Optional[int] + iterable_end = None # type: Optional[int] + for i, rval in enumerate(rvalues): + if isinstance(rval, StarExpr): + typs = get_proper_type(self.expr_checker.visit_star_expr(rval).type) + if self.type_is_iterable(typs) and isinstance(typs, Instance): + if iterable_start is None: + iterable_start = i + iterable_end = i + if (iterable_start is not None + and iterable_end is not None + and iterable_type is not None): + iterable_num = iterable_end - iterable_start + 1 + rvalue_needed = len(lvalues) - (len(rvalues) - iterable_num) + if rvalue_needed > 0: + rvalues = rvalues[0: iterable_start] + [TempNode(iterable_type) + for i in range(rvalue_needed)] + rvalues[iterable_end + 1:] if self.check_rvalue_count_in_assignment(lvalues, len(rvalues), context): star_index = next((i for i, lv in enumerate(lvalues) if diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test index 2ae61f885eae..1c4b537325e8 100644 --- a/test-data/unit/check-tuples.test +++ b/test-data/unit/check-tuples.test @@ -1421,3 +1421,40 @@ t6: Tuple[int, int, int, int, int, int, int, int, int, int, int, int] = (1, 2, 3 # E: Incompatible types in assignment (expression has type Tuple[int, int, ... <15 more items>], variable has type Tuple[int, int, ... <10 more items>]) [builtins fixtures/tuple.pyi] + +[case testTupleWithStarExpr] +from typing import Tuple, List +points = (1, "test") # type: Tuple[int, str] +x, y, z = *points, 0 +reveal_type(x) # N: Revealed type is 'builtins.int' +reveal_type(y) # N: Revealed type is 'builtins.str' +reveal_type(z) # N: Revealed type is 'builtins.int' + +points2 = [1,2] +x2, y2, z2= *points2, "test" + +reveal_type(x2) # N: Revealed type is 'builtins.int*' +reveal_type(y2) # N: Revealed type is 'builtins.int*' +reveal_type(z2) # N: Revealed type is 'builtins.str' + +x3, x4, y3, y4, z3 = *points, *points2, "test" + +reveal_type(x3) # N: Revealed type is 'builtins.int' +reveal_type(x4) # N: Revealed type is 'builtins.str' +reveal_type(y3) # N: Revealed type is 'builtins.int*' +reveal_type(y4) # N: Revealed type is 'builtins.int*' +reveal_type(z3) # N: Revealed type is 'builtins.str' + +x5, x6, y5, y6, z4 = *points2, *points2, "test" + +reveal_type(x5) # N: Revealed type is 'builtins.int*' +reveal_type(x6) # N: Revealed type is 'builtins.int*' +reveal_type(y5) # N: Revealed type is 'builtins.int*' +reveal_type(y6) # N: Revealed type is 'builtins.int*' +reveal_type(z4) # N: Revealed type is 'builtins.str' + +points3 = ["test1", "test2"] +x7, x8, y7, y8 = *points2, *points3 # E: Contiguous iterable with same type expected + +x9, y9, x10, y10, z5 = *points2, 1, *points2 # E: Contiguous iterable with same type expected +[builtins fixtures/tuple.pyi]