Support star expressions in lists, tuples, sets (#1953)

This fixes #1722 and addresses most of #704 (but doesn't close it, because of {**x}).

Author: gvanrossum

mypy/checkexpr.py

@@ -14,7 +14,7 @@
     TupleExpr, DictExpr, FuncExpr, SuperExpr, SliceExpr, Context,
     ListComprehension, GeneratorExpr, SetExpr, MypyFile, Decorator,
     ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
-    DictionaryComprehension, ComplexExpr, EllipsisExpr,
+    DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr,
     TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR2
 )
 from mypy.nodes import function_type
@@ -1286,15 +1286,17 @@ def visit_type_alias_expr(self, alias: TypeAliasExpr) -> Type:
 
     def visit_list_expr(self, e: ListExpr) -> Type:
         """Type check a list expression [...]."""
-        return self.check_list_or_set_expr(e.items, 'builtins.list', '<list>',
-                                           e)
+        return self.check_lst_expr(e.items, 'builtins.list', '<list>', e)
 
     def visit_set_expr(self, e: SetExpr) -> Type:
-        return self.check_list_or_set_expr(e.items, 'builtins.set', '<set>', e)
+        return self.check_lst_expr(e.items, 'builtins.set', '<set>', e)
 
-    def check_list_or_set_expr(self, items: List[Node], fullname: str,
-                               tag: str, context: Context) -> Type:
+    def check_lst_expr(self, items: List[Node], fullname: str,
+                       tag: str, context: Context) -> Type:
         # Translate into type checking a generic function call.
+        # Used for list and set expressions, as well as for tuples
+        # containing star expressions that don't refer to a
+        # Tuple. (Note: "lst" stands for list-set-tuple. :-)
         tvdef = TypeVarDef('T', -1, [], self.chk.object_type())
         tv = TypeVarType(tvdef)
         constructor = CallableType(
@@ -1306,28 +1308,54 @@ def check_list_or_set_expr(self, items: List[Node], fullname: str,
             name=tag,
             variables=[tvdef])
         return self.check_call(constructor,
-                               items,
-                               [nodes.ARG_POS] * len(items), context)[0]
+                               [(i.expr if isinstance(i, StarExpr) else i)
+                                for i in items],
+                               [(nodes.ARG_STAR if isinstance(i, StarExpr) else nodes.ARG_POS)
+                                for i in items],
+                               context)[0]
 
     def visit_tuple_expr(self, e: TupleExpr) -> Type:
         """Type check a tuple expression."""
         ctx = None  # type: TupleType
         # Try to determine type context for type inference.
         if isinstance(self.chk.type_context[-1], TupleType):
             t = self.chk.type_context[-1]
-            if len(t.items) == len(e.items):
-                ctx = t
-        # Infer item types.
+            ctx = t
+        # NOTE: it's possible for the context to have a different
+        # number of items than e.  In that case we use those context
+        # items that match a position in e, and we'll worry about type
+        # mismatches later.
+
+        # Infer item types.  Give up if there's a star expression
+        # that's not a Tuple.
         items = []  # type: List[Type]
+        j = 0  # Index into ctx.items; irrelevant if ctx is None.
         for i in range(len(e.items)):
             item = e.items[i]
             tt = None  # type: Type
-            if not ctx:
-                tt = self.accept(item)
+            if isinstance(item, StarExpr):
+                # Special handling for star expressions.
+                # TODO: If there's a context, and item.expr is a
+                # TupleExpr, flatten it, so we can benefit from the
+                # context?  Counterargument: Why would anyone write
+                # (1, *(2, 3)) instead of (1, 2, 3) except in a test?
+                tt = self.accept(item.expr)
+                self.check_not_void(tt, e)
+                if isinstance(tt, TupleType):
+                    items.extend(tt.items)
+                    j += len(tt.items)
+                else:
+                    # A star expression that's not a Tuple.
+                    # Treat the whole thing as a variable-length tuple.
+                    return self.check_lst_expr(e.items, 'builtins.tuple', '<tuple>', e)
             else:
-                tt = self.accept(item, ctx.items[i])
-            self.check_not_void(tt, e)
-            items.append(tt)
+                if not ctx or j >= len(ctx.items):
+                    tt = self.accept(item)
+                else:
+                    tt = self.accept(item, ctx.items[j])
+                    j += 1
+                self.check_not_void(tt, e)
+                items.append(tt)
         fallback_item = join.join_type_list(items)
         return TupleType(items, self.chk.named_generic_type('builtins.tuple', [fallback_item]))
 

mypy/semanal.py

@@ -1919,14 +1919,20 @@ def visit_super_expr(self, expr: SuperExpr) -> None:
 
     def visit_tuple_expr(self, expr: TupleExpr) -> None:
         for item in expr.items:
+            if isinstance(item, StarExpr):
+                item.valid = True
             item.accept(self)
 
     def visit_list_expr(self, expr: ListExpr) -> None:
         for item in expr.items:
+            if isinstance(item, StarExpr):
+                item.valid = True
             item.accept(self)
 
     def visit_set_expr(self, expr: SetExpr) -> None:
         for item in expr.items:
+            if isinstance(item, StarExpr):
+                item.valid = True
             item.accept(self)
 
     def visit_dict_expr(self, expr: DictExpr) -> None:
@@ -1936,6 +1942,7 @@ def visit_dict_expr(self, expr: DictExpr) -> None:
 
     def visit_star_expr(self, expr: StarExpr) -> None:
         if not expr.valid:
+            # XXX TODO Change this error message
             self.fail('Can use starred expression only as assignment target', expr)
         else:
             expr.expr.accept(self)

test-data/unit/check-inference.test

@@ -809,6 +809,14 @@ s = s_i()
 s = s_s() # E: Incompatible types in assignment (expression has type Set[str], variable has type Set[int])
 [builtins fixtures/set.py]
 
+[case testSetWithStarExpr]
+# options: fast_parser
+s = {1, 2, *(3, 4)}
+t = {1, 2, *s}
+reveal_type(s)  # E: Revealed type is 'builtins.set[builtins.int*]'
+reveal_type(t)  # E: Revealed type is 'builtins.set[builtins.int*]'
+[builtins fixtures/set.py]
+
 
 -- For statements
 -- --------------

test-data/unit/check-lists.test

@@ -63,6 +63,12 @@ class C: pass
 [out]
 main: note: In function "f":
 
-[case testListContainingStarExpr]
-a = [1, *[2]]  # E: Can use starred expression only as assignment target
+[case testListWithStarExpr]
+(x, *a) = [1, 2, 3]
+a = [1, *[2, 3]]
+reveal_type(a)  # E: Revealed type is 'builtins.list[builtins.int]'
+b = [0, *a]
+reveal_type(b)  # E: Revealed type is 'builtins.list[builtins.int*]'
+c = [*a, 0]
+reveal_type(c)  # E: Revealed type is 'builtins.list[builtins.int*]'
 [builtins fixtures/list.py]

test-data/unit/check-tuples.test

@@ -768,3 +768,43 @@ a = ()
 from typing import Sized
 a = None  # type: Sized
 a = ()
+
+[case testTupleWithStarExpr1]
+# options: fast_parser
+a = (1, 2)
+b = (*a, '')
+reveal_type(b)  # E: Revealed type is 'Tuple[builtins.int, builtins.int, builtins.str]'
+
+[case testTupleWithStarExpr2]
+a = [1]
+b = (0, *a)
+reveal_type(b)  # E: Revealed type is 'builtins.tuple[builtins.int*]'
+[builtins fixtures/tuple.py]
+
+[case testTupleWithStarExpr3]
+a = ['']
+b = (0, *a)
+reveal_type(b)  # E: Revealed type is 'builtins.tuple[builtins.object*]'
+c = (*a, '')
+reveal_type(c)  # E: Revealed type is 'builtins.tuple[builtins.str*]'
+[builtins fixtures/tuple.py]
+
+[case testTupleWithStarExpr4]
+a = (1, 1, 'x', 'x')
+b = (1, 'x')
+a = (0, *b, '')
+[builtins fixtures/tuple.py]
+
+[case testTupleWithUndersizedContext]
+a = ([1], 'x')
+a = ([], 'x', 1)  # E: Incompatible types in assignment (expression has type "Tuple[List[int], str, int]", variable has type "Tuple[List[int], str]")
+[builtins fixtures/tuple.py]
+
+[case testTupleWithOversizedContext]
+a = (1, [1], 'x')
+a = (1, [])  # E: Incompatible types in assignment (expression has type "Tuple[int, List[int]]", variable has type "Tuple[int, List[int], str]")
+[builtins fixtures/tuple.py]
+
+[case testTupleWithoutContext]
+a = (1, [])  # E: Need type annotation for variable
+[builtins fixtures/tuple.py]

test-data/unit/fixtures/tuple.py

@@ -1,6 +1,6 @@
 # Builtins stub used in tuple-related test cases.
 
-from typing import Iterable, TypeVar, Generic, Sequence
+from typing import Iterable, Iterator, TypeVar, Generic, Sequence
 
 Tco = TypeVar('Tco', covariant=True)
 
@@ -11,6 +11,7 @@ class type:
     def __init__(self, *a) -> None: pass
     def __call__(self, *a) -> object: pass
 class tuple(Sequence[Tco], Generic[Tco]):
+    def __iter__(self) -> Iterator[Tco]: pass
     def __getitem__(self, x: int) -> Tco: pass
 class function: pass
 
@@ -21,6 +22,8 @@ class str: pass # For convenience
 
 T = TypeVar('T')
 
+class list(Sequence[T], Generic[T]): pass
+
 def sum(iterable: Iterable[T], start: T = None) -> T: pass
 
 True = bool()

test-data/unit/lib-stub/typing.py

@@ -70,7 +70,7 @@ def __aiter__(self) -> 'AsyncIterator[T]': return self
     @abstractmethod
     def __anext__(self) -> Awaitable[T]: pass
 
-class Sequence(Generic[T]):
+class Sequence(Iterable[T], Generic[T]):
     @abstractmethod
     def __getitem__(self, n: Any) -> T: pass
 

test-data/unit/semanal-errors.test

@@ -486,12 +486,8 @@ b = 1
 c = 1
 d = 1
 a = *b
-a = b, (c, *d)
-a = [1, *[2]]
 [out]
 main:4: error: Can use starred expression only as assignment target
-main:5: error: Can use starred expression only as assignment target
-main:6: error: Can use starred expression only as assignment target
 
 [case testStarExpressionInExp]
 a = 1