Tweak constraints handling for splitting typevartuples

mypy/constraints.py

@@ -54,6 +54,7 @@
     extract_unpack,
     find_unpack_in_list,
     split_with_instance,
+    split_with_mapped_and_template,
     split_with_prefix_and_suffix,
 )
 
@@ -677,10 +678,14 @@ 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
-                    )
+                    (
+                        mapped_prefix,
+                        mapped_middle,
+                        mapped_suffix,
+                        template_prefix,
+                        template_middle,
+                        template_suffix,
+                    ) = split_with_mapped_and_template(mapped, template)
 
                     # Add a constraint for the type var tuple, and then
                     # remove it for the case below.

mypy/typevartuples.py

@@ -44,6 +44,44 @@ def split_with_instance(
     )
 
 
+def split_with_mapped_and_template(
+    mapped: Instance, template: Instance
+) -> tuple[
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+    tuple[Type, ...],
+]:
+    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
+    unpack_suffix = len(template_middle) - unpack_prefix - 1
+
+    (
+        mapped_middle_prefix,
+        mapped_middle_middle,
+        mapped_middle_suffix,
+    ) = split_with_prefix_and_suffix(mapped_middle, unpack_prefix, unpack_suffix)
+    (
+        template_middle_prefix,
+        template_middle_middle,
+        template_middle_suffix,
+    ) = split_with_prefix_and_suffix(template_middle, unpack_prefix, unpack_suffix)
+
+    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 extract_unpack(types: Sequence[Type]) -> ProperType | None:
     """Given a list of types, extracts either a single type from an unpack, or returns None."""
     if len(types) == 1:

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

@@ -257,3 +257,50 @@ class Array(Generic[Unpack[Shape]]):
 
 x: Array[float, Height, Width] = Array()
 [builtins fixtures/tuple.pyi]
+
+[case testPep646TypeConcatenation]
+from typing import Generic, TypeVar, NewType
+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 add_batch_axis(x: Array[Unpack[Shape]]) -> Array[Batch, Unpack[Shape]]: ...
+def del_batch_axis(x: Array[Batch, Unpack[Shape]]) -> Array[Unpack[Shape]]: ...
+def add_batch_channels(
+  x: Array[Unpack[Shape]]
+) -> Array[Batch, Unpack[Shape], Channels]: ...
+
+a: Array[Height, Width]
+b = add_batch_axis(a)
+reveal_type(b)  # N: Revealed type is "__main__.Array[__main__.Batch, __main__.Height, __main__.Width]"
+c = del_batch_axis(b)
+reveal_type(c)  # N: Revealed type is "__main__.Array[__main__.Height, __main__.Width]"
+d = add_batch_channels(a)
+reveal_type(d)  # N: Revealed type is "__main__.Array[__main__.Batch, __main__.Height, __main__.Width, __main__.Channels]"
+
+[builtins fixtures/tuple.pyi]
+[case testPep646TypeVarConcatenation]
+from typing import Generic, TypeVar, NewType, Tuple
+from typing_extensions import TypeVarTuple, Unpack
+
+T = TypeVar('T')
+Ts = TypeVarTuple('Ts')
+
+def prefix_tuple(
+    x: T,
+    y: Tuple[Unpack[Ts]],
+) -> Tuple[T, Unpack[Ts]]:
+    ...
+
+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]