Skip to content

Commit d54e8b3

Browse files
authored
Support variadic tuple packing/unpacking (#16205)
This is includes also related things such as tuple concatenation, special-cased tuple "re-packing", and star tuple unpacking in homogeneous collections. It looks like we are very close to the finish line (the only major missing feature is type narrowing using `len()`, apart from this I just need to do couple technical things, and make one final search for missed code paths). Some notes: * Unfortunately, star items on l.h.s create lists at runtime. This means there are various cases where `list[object]` is the best type we can have. * Note I now infer "precise" types for expressions like `(x, *y, z)`, where `y` is say `tuple[int, ...]`. This may cause errors for code that previously worked (when we will turn this feature on). For example `(1, *[], 2)[42]` will be an error. As usual, I propose to try to be strict, and relax if people will complain (FWIW, I expect very few false positives from this). * It may look like `Unpack` can now "leak" if it was never used explicitly. This is not the case, it is just that experimental features are enabled in tests. * There are couple minor changes that affect code without variadic types. Previously tuple type context was used inconsistently for situations with star unpacks, I clean it up a bit (for my tests). Also I infer `Any`-like l.h.s types after an error in tuple unpacking (when needed) to avoid extra "Cannot determine type" errors in my tests.
1 parent 10dfafe commit d54e8b3

File tree

8 files changed

+437
-23
lines changed

8 files changed

+437
-23
lines changed

mypy/argmap.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
Type,
1515
TypedDictType,
1616
TypeOfAny,
17+
TypeVarTupleType,
18+
UnpackType,
1719
get_proper_type,
1820
)
1921

@@ -174,6 +176,7 @@ def expand_actual_type(
174176
actual_kind: nodes.ArgKind,
175177
formal_name: str | None,
176178
formal_kind: nodes.ArgKind,
179+
allow_unpack: bool = False,
177180
) -> Type:
178181
"""Return the actual (caller) type(s) of a formal argument with the given kinds.
179182
@@ -189,6 +192,11 @@ def expand_actual_type(
189192
original_actual = actual_type
190193
actual_type = get_proper_type(actual_type)
191194
if actual_kind == nodes.ARG_STAR:
195+
if isinstance(actual_type, TypeVarTupleType):
196+
# This code path is hit when *Ts is passed to a callable and various
197+
# special-handling didn't catch this. The best thing we can do is to use
198+
# the upper bound.
199+
actual_type = get_proper_type(actual_type.upper_bound)
192200
if isinstance(actual_type, Instance) and actual_type.args:
193201
from mypy.subtypes import is_subtype
194202

@@ -209,7 +217,20 @@ def expand_actual_type(
209217
self.tuple_index = 1
210218
else:
211219
self.tuple_index += 1
212-
return actual_type.items[self.tuple_index - 1]
220+
item = actual_type.items[self.tuple_index - 1]
221+
if isinstance(item, UnpackType) and not allow_unpack:
222+
# An upack item that doesn't have special handling, use upper bound as above.
223+
unpacked = get_proper_type(item.type)
224+
if isinstance(unpacked, TypeVarTupleType):
225+
fallback = get_proper_type(unpacked.upper_bound)
226+
else:
227+
fallback = unpacked
228+
assert (
229+
isinstance(fallback, Instance)
230+
and fallback.type.fullname == "builtins.tuple"
231+
)
232+
item = fallback.args[0]
233+
return item
213234
elif isinstance(actual_type, ParamSpecType):
214235
# ParamSpec is valid in *args but it can't be unpacked.
215236
return actual_type

mypy/checker.py

Lines changed: 99 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -205,10 +205,13 @@
205205
TypeType,
206206
TypeVarId,
207207
TypeVarLikeType,
208+
TypeVarTupleType,
208209
TypeVarType,
209210
UnboundType,
210211
UninhabitedType,
211212
UnionType,
213+
UnpackType,
214+
find_unpack_in_list,
212215
flatten_nested_unions,
213216
get_proper_type,
214217
get_proper_types,
@@ -3430,6 +3433,37 @@ def is_assignable_slot(self, lvalue: Lvalue, typ: Type | None) -> bool:
34303433
return all(self.is_assignable_slot(lvalue, u) for u in typ.items)
34313434
return False
34323435

3436+
def flatten_rvalues(self, rvalues: list[Expression]) -> list[Expression]:
3437+
"""Flatten expression list by expanding those * items that have tuple type.
3438+
3439+
For each regular type item in the tuple type use a TempNode(), for an Unpack
3440+
item use a corresponding StarExpr(TempNode()).
3441+
"""
3442+
new_rvalues = []
3443+
for rv in rvalues:
3444+
if not isinstance(rv, StarExpr):
3445+
new_rvalues.append(rv)
3446+
continue
3447+
typ = get_proper_type(self.expr_checker.accept(rv.expr))
3448+
if not isinstance(typ, TupleType):
3449+
new_rvalues.append(rv)
3450+
continue
3451+
for t in typ.items:
3452+
if not isinstance(t, UnpackType):
3453+
new_rvalues.append(TempNode(t))
3454+
else:
3455+
unpacked = get_proper_type(t.type)
3456+
if isinstance(unpacked, TypeVarTupleType):
3457+
fallback = unpacked.upper_bound
3458+
else:
3459+
assert (
3460+
isinstance(unpacked, Instance)
3461+
and unpacked.type.fullname == "builtins.tuple"
3462+
)
3463+
fallback = unpacked
3464+
new_rvalues.append(StarExpr(TempNode(fallback)))
3465+
return new_rvalues
3466+
34333467
def check_assignment_to_multiple_lvalues(
34343468
self,
34353469
lvalues: list[Lvalue],
@@ -3439,18 +3473,16 @@ def check_assignment_to_multiple_lvalues(
34393473
) -> None:
34403474
if isinstance(rvalue, (TupleExpr, ListExpr)):
34413475
# Recursively go into Tuple or List expression rhs instead of
3442-
# using the type of rhs, because this allowed more fine grained
3476+
# using the type of rhs, because this allows more fine-grained
34433477
# control in cases like: a, b = [int, str] where rhs would get
34443478
# type List[object]
34453479
rvalues: list[Expression] = []
34463480
iterable_type: Type | None = None
34473481
last_idx: int | None = None
3448-
for idx_rval, rval in enumerate(rvalue.items):
3482+
for idx_rval, rval in enumerate(self.flatten_rvalues(rvalue.items)):
34493483
if isinstance(rval, StarExpr):
34503484
typs = get_proper_type(self.expr_checker.accept(rval.expr))
3451-
if isinstance(typs, TupleType):
3452-
rvalues.extend([TempNode(typ) for typ in typs.items])
3453-
elif self.type_is_iterable(typs) and isinstance(typs, Instance):
3485+
if self.type_is_iterable(typs) and isinstance(typs, Instance):
34543486
if iterable_type is not None and iterable_type != self.iterable_item_type(
34553487
typs, rvalue
34563488
):
@@ -3517,8 +3549,32 @@ def check_assignment_to_multiple_lvalues(
35173549
self.check_multi_assignment(lvalues, rvalue, context, infer_lvalue_type)
35183550

35193551
def check_rvalue_count_in_assignment(
3520-
self, lvalues: list[Lvalue], rvalue_count: int, context: Context
3552+
self,
3553+
lvalues: list[Lvalue],
3554+
rvalue_count: int,
3555+
context: Context,
3556+
rvalue_unpack: int | None = None,
35213557
) -> bool:
3558+
if rvalue_unpack is not None:
3559+
if not any(isinstance(e, StarExpr) for e in lvalues):
3560+
self.fail("Variadic tuple unpacking requires a star target", context)
3561+
return False
3562+
if len(lvalues) > rvalue_count:
3563+
self.fail(message_registry.TOO_MANY_TARGETS_FOR_VARIADIC_UNPACK, context)
3564+
return False
3565+
left_star_index = next(i for i, lv in enumerate(lvalues) if isinstance(lv, StarExpr))
3566+
left_prefix = left_star_index
3567+
left_suffix = len(lvalues) - left_star_index - 1
3568+
right_prefix = rvalue_unpack
3569+
right_suffix = rvalue_count - rvalue_unpack - 1
3570+
if left_suffix > right_suffix or left_prefix > right_prefix:
3571+
# Case of asymmetric unpack like:
3572+
# rv: tuple[int, *Ts, int, int]
3573+
# x, y, *xs, z = rv
3574+
# it is technically valid, but is tricky to reason about.
3575+
# TODO: support this (at least if the r.h.s. unpack is a homogeneous tuple).
3576+
self.fail(message_registry.TOO_MANY_TARGETS_FOR_VARIADIC_UNPACK, context)
3577+
return True
35223578
if any(isinstance(lvalue, StarExpr) for lvalue in lvalues):
35233579
if len(lvalues) - 1 > rvalue_count:
35243580
self.msg.wrong_number_values_to_unpack(rvalue_count, len(lvalues) - 1, context)
@@ -3552,6 +3608,13 @@ def check_multi_assignment(
35523608
if len(relevant_items) == 1:
35533609
rvalue_type = get_proper_type(relevant_items[0])
35543610

3611+
if (
3612+
isinstance(rvalue_type, TupleType)
3613+
and find_unpack_in_list(rvalue_type.items) is not None
3614+
):
3615+
# Normalize for consistent handling with "old-style" homogeneous tuples.
3616+
rvalue_type = expand_type(rvalue_type, {})
3617+
35553618
if isinstance(rvalue_type, AnyType):
35563619
for lv in lvalues:
35573620
if isinstance(lv, StarExpr):
@@ -3663,7 +3726,10 @@ def check_multi_assignment_from_tuple(
36633726
undefined_rvalue: bool,
36643727
infer_lvalue_type: bool = True,
36653728
) -> None:
3666-
if self.check_rvalue_count_in_assignment(lvalues, len(rvalue_type.items), context):
3729+
rvalue_unpack = find_unpack_in_list(rvalue_type.items)
3730+
if self.check_rvalue_count_in_assignment(
3731+
lvalues, len(rvalue_type.items), context, rvalue_unpack=rvalue_unpack
3732+
):
36673733
star_index = next(
36683734
(i for i, lv in enumerate(lvalues) if isinstance(lv, StarExpr)), len(lvalues)
36693735
)
@@ -3708,12 +3774,37 @@ def check_multi_assignment_from_tuple(
37083774
self.check_assignment(lv, self.temp_node(rv_type, context), infer_lvalue_type)
37093775
if star_lv:
37103776
list_expr = ListExpr(
3711-
[self.temp_node(rv_type, context) for rv_type in star_rv_types]
3777+
[
3778+
self.temp_node(rv_type, context)
3779+
if not isinstance(rv_type, UnpackType)
3780+
else StarExpr(self.temp_node(rv_type.type, context))
3781+
for rv_type in star_rv_types
3782+
]
37123783
)
37133784
list_expr.set_line(context)
37143785
self.check_assignment(star_lv.expr, list_expr, infer_lvalue_type)
37153786
for lv, rv_type in zip(right_lvs, right_rv_types):
37163787
self.check_assignment(lv, self.temp_node(rv_type, context), infer_lvalue_type)
3788+
else:
3789+
# Store meaningful Any types for lvalues, errors are already given
3790+
# by check_rvalue_count_in_assignment()
3791+
if infer_lvalue_type:
3792+
for lv in lvalues:
3793+
if (
3794+
isinstance(lv, NameExpr)
3795+
and isinstance(lv.node, Var)
3796+
and lv.node.type is None
3797+
):
3798+
lv.node.type = AnyType(TypeOfAny.from_error)
3799+
elif isinstance(lv, StarExpr):
3800+
if (
3801+
isinstance(lv.expr, NameExpr)
3802+
and isinstance(lv.expr.node, Var)
3803+
and lv.expr.node.type is None
3804+
):
3805+
lv.expr.node.type = self.named_generic_type(
3806+
"builtins.list", [AnyType(TypeOfAny.from_error)]
3807+
)
37173808

37183809
def lvalue_type_for_inference(self, lvalues: list[Lvalue], rvalue_type: TupleType) -> Type:
37193810
star_index = next(

0 commit comments

Comments
 (0)