Lenient handling of *tuple[Any, ...] (part 2) (#22006)

Ref #19858

This (smaller) part handles various edge cases that should be handled
leniently in presence of `*tuple[Any, ...]`: indexing and unpacking.
Everything is straightforward, similar to first part
#22001, the tests are focused on
situations that previously gave errors.

Author: ilevkivskyi

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)
             )

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)

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

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

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]

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