Extend support for tuple subtyping with typevar tuples (#13718)

This adds a new Pep646 test case which demonstrates that like with the
constraints from PR #13716 we need to split twice for subtyping when
handling Unpacks.

It also demonstrates a weakness of the previous PR which is that the
middle-prefix and the prefix may need to be handled differently so we
introduce another splitting function that returns a 10-tuple instead of
a 6-tuple and reimplement the 6-tuple version on top of the 10-tuple
version.

Complicating things further, the test case reveals that there are error
cases where split_with_mapped_and_template cannot actually unpack the
middle a second time because the mapped middle is too short to do the
unpack. We also now have to deal with the case where there was no unpack
in the template in which case we only do a single split.

In addition we fix a behavioral issue where according to PEP646 we
should assume that Tuple[Unpack[Tuple[Any, ...]]] is equivalent to
Tuple[Any, Any] even if we don't actually know the lengths match. As
such test_type_var_tuple_unpacked_variable_length_tuple changes from
asserting a strict subtype to asserting equivalence.

One of the messages was bad as well so we add a branch for UnpackType in
message pretty-printing.

Author: jhance

mypy/constraints.py

@@ -678,14 +678,20 @@ def visit_instance(self, template: Instance) -> list[Constraint]:
                 mapped = map_instance_to_supertype(instance, template.type)
                 tvars = template.type.defn.type_vars
                 if template.type.has_type_var_tuple_type:
+                    mapped_prefix, mapped_middle, mapped_suffix = split_with_instance(mapped)
+                    template_prefix, template_middle, template_suffix = split_with_instance(
+                        template
+                    )
+                    split_result = split_with_mapped_and_template(mapped, template)
+                    assert split_result is not None
                     (
                         mapped_prefix,
                         mapped_middle,
                         mapped_suffix,
                         template_prefix,
                         template_middle,
                         template_suffix,
-                    ) = split_with_mapped_and_template(mapped, template)
+                    ) = split_result
 
                     # Add a constraint for the type var tuple, and then
                     # remove it for the case below.

mypy/messages.py

@@ -84,6 +84,7 @@
     UnboundType,
     UninhabitedType,
     UnionType,
+    UnpackType,
     get_proper_type,
     get_proper_types,
 )
@@ -2257,6 +2258,8 @@ def format_literal_value(typ: LiteralType) -> str:
         else:
             # There are type arguments. Convert the arguments to strings.
             return f"{base_str}[{format_list(itype.args)}]"
+    elif isinstance(typ, UnpackType):
+        return f"Unpack[{format(typ.type)}]"
     elif isinstance(typ, TypeVarType):
         # This is similar to non-generic instance types.
         return typ.name

mypy/subtypes.py

@@ -62,7 +62,7 @@
 )
 from mypy.typestate import SubtypeKind, TypeState
 from mypy.typevars import fill_typevars_with_any
-from mypy.typevartuples import extract_unpack, split_with_instance
+from mypy.typevartuples import extract_unpack, fully_split_with_mapped_and_template
 
 # Flags for detected protocol members
 IS_SETTABLE: Final = 1
@@ -485,8 +485,22 @@ def visit_instance(self, left: Instance) -> bool:
                     t = erased
                 nominal = True
                 if right.type.has_type_var_tuple_type:
-                    left_prefix, left_middle, left_suffix = split_with_instance(left)
-                    right_prefix, right_middle, right_suffix = split_with_instance(right)
+                    split_result = fully_split_with_mapped_and_template(left, right)
+                    if split_result is None:
+                        return False
+
+                    (
+                        left_prefix,
+                        left_mprefix,
+                        left_middle,
+                        left_msuffix,
+                        left_suffix,
+                        right_prefix,
+                        right_mprefix,
+                        right_middle,
+                        right_msuffix,
+                        right_suffix,
+                    ) = split_result
 
                     left_unpacked = extract_unpack(left_middle)
                     right_unpacked = extract_unpack(right_middle)
@@ -495,6 +509,15 @@ def visit_instance(self, left: Instance) -> bool:
                     def check_mixed(
                         unpacked_type: ProperType, compare_to: tuple[Type, ...]
                     ) -> bool:
+                        if (
+                            isinstance(unpacked_type, Instance)
+                            and unpacked_type.type.fullname == "builtins.tuple"
+                        ):
+                            if not all(
+                                is_equivalent(l, unpacked_type.args[0]) for l in compare_to
+                            ):
+                                return False
+                            return True
                         if isinstance(unpacked_type, TypeVarTupleType):
                             return False
                         if isinstance(unpacked_type, AnyType):
@@ -521,13 +544,6 @@ def check_mixed(
                         if not check_mixed(left_unpacked, right_middle):
                             return False
                     elif left_unpacked is None and right_unpacked is not None:
-                        if (
-                            isinstance(right_unpacked, Instance)
-                            and right_unpacked.type.fullname == "builtins.tuple"
-                        ):
-                            return all(
-                                is_equivalent(l, right_unpacked.args[0]) for l in left_middle
-                            )
                         if not check_mixed(right_unpacked, left_middle):
                             return False
 
@@ -540,16 +556,24 @@ def check_mixed(
                             if not is_equivalent(left_t, right_t):
                                 return False
 
+                    assert len(left_mprefix) == len(right_mprefix)
+                    assert len(left_msuffix) == len(right_msuffix)
+
+                    for left_item, right_item in zip(
+                        left_mprefix + left_msuffix, right_mprefix + right_msuffix
+                    ):
+                        if not is_equivalent(left_item, right_item):
+                            return False
+
                     left_items = t.args[: right.type.type_var_tuple_prefix]
                     right_items = right.args[: right.type.type_var_tuple_prefix]
                     if right.type.type_var_tuple_suffix:
                         left_items += t.args[-right.type.type_var_tuple_suffix :]
                         right_items += right.args[-right.type.type_var_tuple_suffix :]
-
                     unpack_index = right.type.type_var_tuple_prefix
                     assert unpack_index is not None
                     type_params = zip(
-                        left_prefix + right_suffix,
+                        left_prefix + left_suffix,
                         right_prefix + right_suffix,
                         right.type.defn.type_vars[:unpack_index]
                         + right.type.defn.type_vars[unpack_index + 1 :],

mypy/test/testsubtypes.py

@@ -273,6 +273,22 @@ def test_type_var_tuple_with_prefix_suffix(self) -> None:
             Instance(self.fx.gvi, [self.fx.a, UnpackType(self.fx.ss), self.fx.b, self.fx.c]),
         )
 
+    def test_type_var_tuple_unpacked_varlength_tuple(self) -> None:
+        self.assert_subtype(
+            Instance(
+                self.fx.gvi,
+                [
+                    UnpackType(
+                        TupleType(
+                            [self.fx.a, self.fx.b],
+                            fallback=Instance(self.fx.std_tuplei, [self.fx.o]),
+                        )
+                    )
+                ],
+            ),
+            Instance(self.fx.gvi, [self.fx.a, self.fx.b]),
+        )
+
     def test_type_var_tuple_unpacked_tuple(self) -> None:
         self.assert_subtype(
             Instance(
@@ -333,7 +349,7 @@ def test_type_var_tuple_unpacked_tuple(self) -> None:
         )
 
     def test_type_var_tuple_unpacked_variable_length_tuple(self) -> None:
-        self.assert_strict_subtype(
+        self.assert_equivalent(
             Instance(self.fx.gvi, [self.fx.a, self.fx.a]),
             Instance(self.fx.gvi, [UnpackType(Instance(self.fx.std_tuplei, [self.fx.a]))]),
         )

mypy/typevartuples.py

@@ -53,13 +53,70 @@ def split_with_mapped_and_template(
     tuple[Type, ...],
     tuple[Type, ...],
     tuple[Type, ...],
-]:
+] | None:
+    split_result = fully_split_with_mapped_and_template(mapped, template)
+    if split_result is None:
+        return None
+
+    (
+        mapped_prefix,
+        mapped_middle_prefix,
+        mapped_middle_middle,
+        mapped_middle_suffix,
+        mapped_suffix,
+        template_prefix,
+        template_middle_prefix,
+        template_middle_middle,
+        template_middle_suffix,
+        template_suffix,
+    ) = split_result
+
+    return (
+        mapped_prefix + mapped_middle_prefix,
+        mapped_middle_middle,
+        mapped_middle_suffix + mapped_suffix,
+        template_prefix + template_middle_prefix,
+        template_middle_middle,
+        template_middle_suffix + template_suffix,
+    )
+
+
+def fully_split_with_mapped_and_template(
+    mapped: Instance, template: Instance
+) -> tuple[
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+] | None:
     mapped_prefix, mapped_middle, mapped_suffix = split_with_instance(mapped)
     template_prefix, template_middle, template_suffix = split_with_instance(template)
 
     unpack_prefix = find_unpack_in_list(template_middle)
-    assert unpack_prefix is not None
+    if unpack_prefix is None:
+        return (
+            mapped_prefix,
+            (),
+            mapped_middle,
+            (),
+            mapped_suffix,
+            template_prefix,
+            (),
+            template_middle,
+            (),
+            template_suffix,
+        )
+
     unpack_suffix = len(template_middle) - unpack_prefix - 1
+    # mapped_middle is too short to do the unpack
+    if unpack_prefix + unpack_suffix > len(mapped_middle):
+        return None
 
     (
         mapped_middle_prefix,
@@ -73,12 +130,16 @@ def split_with_mapped_and_template(
     ) = split_with_prefix_and_suffix(template_middle, unpack_prefix, unpack_suffix)
 
     return (
-        mapped_prefix + mapped_middle_prefix,
+        mapped_prefix,
+        mapped_middle_prefix,
         mapped_middle_middle,
-        mapped_middle_suffix + mapped_suffix,
-        template_prefix + template_middle_prefix,
+        mapped_middle_suffix,
+        mapped_suffix,
+        template_prefix,
+        template_middle_prefix,
         template_middle_middle,
-        template_middle_suffix + template_suffix,
+        template_middle_suffix,
+        template_suffix,
     )
 
 

test-data/unit/check-typevar-tuple.test

@@ -304,3 +304,46 @@ def prefix_tuple(
 z = prefix_tuple(x=0, y=(True, 'a'))
 reveal_type(z)  # N: Revealed type is "Tuple[builtins.int, builtins.bool, builtins.str]"
 [builtins fixtures/tuple.pyi]
+[case testPep646TypeVarTupleUnpacking]
+from typing import Generic, TypeVar, NewType, Any, Tuple
+from typing_extensions import TypeVarTuple, Unpack
+
+Shape = TypeVarTuple('Shape')
+
+Channels = NewType("Channels", int)
+Batch = NewType("Batch", int)
+Height = NewType('Height', int)
+Width = NewType('Width', int)
+
+class Array(Generic[Unpack[Shape]]):
+    pass
+
+def process_batch_channels(
+    x: Array[Batch, Unpack[Tuple[Any, ...]], Channels]
+) -> None:
+    ...
+
+x: Array[Batch, Height, Width, Channels]
+process_batch_channels(x)
+y: Array[Batch, Channels]
+process_batch_channels(y)
+z: Array[Batch]
+process_batch_channels(z)  # E: Argument 1 to "process_batch_channels" has incompatible type "Array[Batch]"; expected "Array[Batch, Unpack[Tuple[Any, ...]], Channels]"
+
+u: Array[Unpack[Tuple[Any, ...]]]
+
+def expect_variadic_array(
+    x: Array[Batch, Unpack[Shape]]
+) -> None:
+    ...
+
+def expect_variadic_array_2(
+    x: Array[Batch, Height, Width, Channels]
+) -> None:
+    ...
+
+expect_variadic_array(u)
+expect_variadic_array_2(u)
+
+
+[builtins fixtures/tuple.pyi]