diff --git a/mypy/checker.py b/mypy/checker.py index d4ba6c54fb35..f03d06d965c3 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -291,10 +291,12 @@ def __init__(self) -> None: UninhabitedType, UnionType, UnpackType, + extend_args_for_prefix_and_suffix, find_unpack_in_list, flatten_nested_unions, get_proper_type, get_proper_types, + get_variadic_item, instance_cache, is_literal_type, is_named_instance, @@ -4389,6 +4391,36 @@ def flatten_lvalues(self, lvalues: list[Expression]) -> list[Expression]: res.append(lv) return res + def adjust_rvalue_type_if_possible( + self, rvalue_type: TupleType, lvalues: list[Lvalue] + ) -> TupleType: + """Adjust type of rvalue to match the shape/structure of lvalues. + + Currently, we only allow this if the rvalue type has contains *tuple[Any, ...]. + """ + right_variadic = get_variadic_item(rvalue_type) + if right_variadic is None: + return rvalue_type + right_unpack_index, right_item = right_variadic + if not isinstance(get_proper_type(right_item), AnyType): + return rvalue_type + left_star_index = next( + (i for i, lv in enumerate(lvalues) if isinstance(lv, StarExpr)), None + ) + if left_star_index is None: + extra = len(lvalues) - len(rvalue_type.items) + 1 + if extra < 0: + return rvalue_type + return rvalue_type.copy_modified( + items=rvalue_type.items[:right_unpack_index] + + [right_item] * extra + + rvalue_type.items[right_unpack_index + 1 :] + ) + new_items = extend_args_for_prefix_and_suffix( + tuple(rvalue_type.items), left_star_index, len(lvalues) - left_star_index - 1 + ) + return rvalue_type.copy_modified(items=list(new_items)) + def check_multi_assignment_from_tuple( self, lvalues: list[Lvalue], @@ -4399,6 +4431,9 @@ def check_multi_assignment_from_tuple( infer_lvalue_type: bool = True, ) -> None: rvalue_unpack = find_unpack_in_list(rvalue_type.items) + if rvalue_unpack is not None: + rvalue_type = self.adjust_rvalue_type_if_possible(rvalue_type, lvalues) + rvalue_unpack = find_unpack_in_list(rvalue_type.items) if self.check_rvalue_count_in_assignment( lvalues, len(rvalue_type.items), context, rvalue_unpack=rvalue_unpack ): @@ -4440,8 +4475,14 @@ def check_multi_assignment_from_tuple( if isinstance(reinferred_rvalue_type, TupleType): # This branch will usually be taken, but in some cases context can # e.g. select a different overload + # TODO: reinferred tuple may be of a different (invalid) shape. rvalue_type = reinferred_rvalue_type + # Reinferring the type can undo the shape adjustment, so do it again. + rvalue_unpack = find_unpack_in_list(rvalue_type.items) + if rvalue_unpack is not None: + rvalue_type = self.adjust_rvalue_type_if_possible(rvalue_type, lvalues) + left_rv_types, star_rv_types, right_rv_types = self.split_around_star( rvalue_type.items, star_index, len(lvalues) ) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index d793797a445d..f8b2fc4e7971 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -200,6 +200,7 @@ flatten_nested_unions, get_proper_type, get_proper_types, + get_variadic_item, has_recursive_types, has_type_vars, is_named_instance, @@ -4654,6 +4655,10 @@ def visit_tuple_index_helper(self, left: TupleType, n: int) -> Type | None: if n >= self.min_tuple_length(left): # For tuple[int, *tuple[str, ...], int] we allow either index 0 or 1, # since variadic item may have zero items. + if isinstance(get_proper_type(middle), AnyType): + # The only exception is when the variadic item is Any, + # which is handled leniently. + return UnionType.make_union([middle] + left.items[unpack_index + 1 :]) return None if n < unpack_index: return left.items[n] @@ -4666,6 +4671,8 @@ def visit_tuple_index_helper(self, left: TupleType, n: int) -> Type | None: n += self.min_tuple_length(left) if n < 0: # Similar to above, we only allow -1, and -2 for tuple[int, *tuple[str, ...], int] + if isinstance(get_proper_type(middle), AnyType): + return UnionType.make_union(left.items[:unpack_index] + [middle]) return None if n >= unpack_index + extra_items: return left.items[n - extra_items + 1] @@ -4698,8 +4705,18 @@ def visit_tuple_slice_helper(self, left_type: TupleType, slic: SliceExpr) -> Typ items: list[Type] = [] for b, e, s in itertools.product(begin, end, stride): + if s == 0: + self.chk.fail("Slice step cannot be zero", slic) + items.append(self.named_type("builtins.tuple")) + continue item = left_type.slice(b, e, s, fallback=self.named_type("builtins.tuple")) if item is None: + left_variadic = get_variadic_item(left_type) + if left_variadic is not None: + _, left_item = left_variadic + if isinstance(get_proper_type(left_item), AnyType): + # If the tuple has *tuple[Any, ...] slice should never fail. + return self.nonliteral_tuple_index_helper(left_type, slic) self.chk.fail(message_registry.AMBIGUOUS_SLICE_OF_VARIADIC_TUPLE, slic) return AnyType(TypeOfAny.from_error) items.append(item) diff --git a/mypy/subtypes.py b/mypy/subtypes.py index 41df445d592f..9c67ebbdf106 100644 --- a/mypy/subtypes.py +++ b/mypy/subtypes.py @@ -73,6 +73,7 @@ find_unpack_in_list, flatten_nested_unions, get_proper_type, + get_variadic_item, is_named_instance, split_with_prefix_and_suffix, ) @@ -852,7 +853,7 @@ def adjust_left_if_possible(self, left: TupleType, right: TupleType) -> TupleTyp Note: this only works if right is fixed size (including *Ts), the variadic right are handled by the caller, currently with variadic_tuple_subtype(). """ - left_variadic = self.get_variadic_item(left) + left_variadic = get_variadic_item(left) if left_variadic is None: return left left_unpack_index, left_item = left_variadic @@ -884,19 +885,6 @@ def adjust_left_if_possible(self, left: TupleType, right: TupleType) -> TupleTyp ) return left.copy_modified(items=list(new_items)) - def get_variadic_item(self, tup: TupleType) -> tuple[int, Type] | None: - """If this is tuple[X, *tuple[Y, ...], Z], return Y, otherwise None.""" - unpack_index = find_unpack_in_list(tup.items) - if unpack_index is None: - return None - unpack = tup.items[unpack_index] - assert isinstance(unpack, UnpackType) - unpacked = get_proper_type(unpack.type) - if not isinstance(unpacked, Instance): - return None - assert unpacked.type.fullname == "builtins.tuple" - return unpack_index, unpacked.args[0] - def variadic_tuple_subtype(self, left: TupleType, right: TupleType) -> bool: """Check subtyping between two potentially variadic tuples. @@ -906,7 +894,7 @@ def variadic_tuple_subtype(self, left: TupleType, right: TupleType) -> bool: Note: the cases where right is fixed or has *Ts unpack should be handled by the caller. """ - right_variadic = self.get_variadic_item(right) + right_variadic = get_variadic_item(right) if right_variadic is None: # This case should be handled by the caller. return False diff --git a/mypy/types.py b/mypy/types.py index 7a1470964d25..e01a1f21e8fb 100644 --- a/mypy/types.py +++ b/mypy/types.py @@ -2946,9 +2946,6 @@ def slice( if fallback is None: fallback = self.partial_fallback - if stride == 0: - return None - if any(isinstance(t, UnpackType) for t in self.items): total = len(self.items) unpack_index = find_unpack_in_list(self.items) @@ -2989,7 +2986,7 @@ def slice( else: return None else: - # TODO: there some additional cases we can support for homogeneous variadic + # TODO: there are some additional cases we can support for homogeneous variadic # items, we can "eat away" finite number of items. return None else: @@ -4448,6 +4445,20 @@ def type_vars_as_args(type_vars: Sequence[TypeVarLikeType]) -> tuple[Type, ...]: return tuple(args) +def get_variadic_item(tup: TupleType) -> tuple[int, Type] | None: + """If this is tuple[X, *tuple[Y, ...], Z], return Y, otherwise None.""" + unpack_index = find_unpack_in_list(tup.items) + if unpack_index is None: + return None + unpack = tup.items[unpack_index] + assert isinstance(unpack, UnpackType) + unpacked = get_proper_type(unpack.type) + if not isinstance(unpacked, Instance): + return None + assert unpacked.type.fullname == "builtins.tuple" + return unpack_index, unpacked.args[0] + + # See docstring for mypy/cache.py for reserved tag ranges. # Instance-related tags. INSTANCE: Final[Tag] = 80 diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test index 79c6f630f4d9..6d523d375dd2 100644 --- a/test-data/unit/check-tuples.test +++ b/test-data/unit/check-tuples.test @@ -1485,8 +1485,7 @@ t[y:] # E: Invalid index type "slice[str, None, None]" for "tuple[int, str]"; e [case testTupleSliceStepZeroNoCrash] # This was crashing: https://github.com/python/mypy/issues/18062 -# TODO: emit better error when 0 is used for step -()[::0] # E: Ambiguous slice of a variadic tuple +()[::0] # E: Slice step cannot be zero [builtins fixtures/tuple.pyi] [case testInferTupleTypeFallbackAgainstInstance] diff --git a/test-data/unit/check-typevar-tuple.test b/test-data/unit/check-typevar-tuple.test index a8898a8329be..eab26463ebf6 100644 --- a/test-data/unit/check-typevar-tuple.test +++ b/test-data/unit/check-typevar-tuple.test @@ -1563,6 +1563,57 @@ x = c1 x = c2 [builtins fixtures/tuple.pyi] +[case testAnyVariadicTupleLenientIndex] +from typing import Any, Unpack + +t: tuple[int, Unpack[tuple[Any, ...]], str] +t1: tuple[int, Unpack[tuple[Any, ...]]] +t2: tuple[Unpack[tuple[Any, ...]], str] + +reveal_type(t[42]) # N: Revealed type is "Any | builtins.str" +reveal_type(t[-42]) # N: Revealed type is "builtins.int | Any" + +reveal_type(t1[42]) # N: Revealed type is "Any" +reveal_type(t1[-42]) # N: Revealed type is "builtins.int | Any" + +reveal_type(t2[42]) # N: Revealed type is "Any | builtins.str" +reveal_type(t2[-42]) # N: Revealed type is "Any" + +# Slicing support for variadic tuples is currently best effort, the key idea is to not give an error. +reveal_type(t[1:3:2]) # N: Revealed type is "builtins.tuple[builtins.int | Any | builtins.str, ...]" +reveal_type(t1[1:3:2]) # N: Revealed type is "builtins.tuple[builtins.int | Any, ...]" +reveal_type(t2[1:3:2]) # N: Revealed type is "builtins.tuple[Any | builtins.str, ...]" +bt = t[::0] # E: Slice step cannot be zero +reveal_type(bt) # N: Revealed type is "builtins.tuple[Any, ...]" +[builtins fixtures/tuple.pyi] + +[case testAnyVariadicTupleLenientUnpacking] +from typing import Any, Unpack + +t: tuple[int, Unpack[tuple[Any, ...]], str] + +x0, y0 = t +reveal_type(x0) # N: Revealed type is "builtins.int" +reveal_type(y0) # N: Revealed type is "builtins.str" + +x, y, z = t +reveal_type(x) # N: Revealed type is "builtins.int" +reveal_type(y) # N: Revealed type is "Any" +reveal_type(z) # N: Revealed type is "builtins.str" + +x1, y1, *zz, x2, y2 = t +reveal_type(x1) # N: Revealed type is "builtins.int" +reveal_type(y1) # N: Revealed type is "Any" +reveal_type(zz) # N: Revealed type is "builtins.list[Any]" +reveal_type(x2) # N: Revealed type is "Any" +reveal_type(y2) # N: Revealed type is "builtins.str" + +t2: tuple[int, int, Unpack[tuple[Any, ...]], str, str] + +# Not the best error, but this definitely invalid, even with Any. +x3, y3, z3 = t2 # E: Variadic tuple unpacking requires a star target +[builtins fixtures/tuple.pyi] + [case testUnpackingVariadicTuplesTypeVar] from typing import Tuple from typing_extensions import TypeVarTuple, Unpack