Lenient handling of *tuple[Any, ...] (part 3) (#22014)

Fixes #20771 (it already worked,
but I am now adding a regression test)
Fixes #20651

This handles various edge cases related to inference with variadic
unpacks (including `Any`). Interestingly, after properly handling some
missing edge cases, I realized that the variadic case can be "merged in
the main branch" of the inference logic, so that the final diff is very
simple.

Author: ilevkivskyi

mypy/checkpattern.py

@@ -402,7 +402,7 @@ def visit_sequence_pattern(self, o: SequencePattern) -> PatternType:
         return PatternType(new_type, rest_type, captures)
 
     def contract_starred_pattern_types(
-        self, types: list[Type], star_pos: int | None, num_patterns: int
+        self, types: list[Type], star_pos: int | None, num_required_patterns: int
     ) -> list[Type]:
         """
         Contracts a list of types in a sequence pattern depending on the position of a starred
@@ -422,15 +422,13 @@ def contract_starred_pattern_types(
             # This should be guaranteed by the normalization in the caller.
             assert isinstance(unpacked, Instance) and unpacked.type.fullname == "builtins.tuple"
             if star_pos is None:
-                missing = num_patterns - len(types) + 1
+                missing = num_required_patterns - len(types) + 1
                 new_types = types[:unpack_index]
                 new_types += [unpacked.args[0]] * missing
                 new_types += types[unpack_index + 1 :]
                 return new_types
             prefix, middle, suffix = split_with_prefix_and_suffix(
-                tuple([UnpackType(unpacked) if isinstance(t, UnpackType) else t for t in types]),
-                star_pos,
-                num_patterns - star_pos,
+                tuple(types), star_pos, num_required_patterns - star_pos
             )
             new_middle = []
             for m in middle:
@@ -445,7 +443,7 @@ def contract_starred_pattern_types(
             if star_pos is None:
                 return types
             new_types = types[:star_pos]
-            star_length = len(types) - num_patterns
+            star_length = len(types) - num_required_patterns
             new_types.append(make_simplified_union(types[star_pos : star_pos + star_length]))
             new_types += types[star_pos + star_length :]
             return new_types

mypy/constraints.py

@@ -1318,7 +1318,7 @@ def visit_tuple_type(self, template: TupleType) -> list[Constraint]:
                     a_unpack = actual.items[a_unpack_index]
                     assert isinstance(a_unpack, UnpackType)
                     a_unpacked = get_proper_type(a_unpack.type)
-                    if len(actual.items) + 1 <= len(template.items):
+                    if len(actual.items) <= len(template.items) + 1:
                         a_prefix_len = a_unpack_index
                         a_suffix_len = len(actual.items) - a_unpack_index - 1
                         t_prefix, t_middle, t_suffix = split_with_prefix_and_suffix(
@@ -1591,6 +1591,17 @@ def build_constraints_for_simple_unpack(
             # This is the only case where we can guarantee there will be no partial overlap
             # (note however partial overlap is OK for variadic tuples, it is handled below).
             t_unpack = template_args[template_unpack]
+        else:
+            # A special case for a variadic actual tuple unpack, we can infer T <: X from
+            # tuple[..., *tuple[T, ...], ...] <: tuple[..., *tuple[X, ...], ...] etc.
+            actual_unpack_type = actual_args[actual_unpack]
+            assert isinstance(actual_unpack_type, UnpackType)
+            a_unpacked = get_proper_type(actual_unpack_type.type)
+            if isinstance(a_unpacked, Instance) and a_unpacked.type.fullname == "builtins.tuple":
+                t_unpack = template_args[template_unpack]
+                # In this case we can "eat away" as much as we need.
+                common_prefix = template_prefix
+                common_suffix = template_suffix
 
     # Handle constraints from prefixes/suffixes first.
     start, middle, end = split_with_prefix_and_suffix(
@@ -1619,18 +1630,6 @@ def build_constraints_for_simple_unpack(
                         res.extend(infer_constraints(tp.args[0], a_tp.args[0], direction))
         elif isinstance(tp, TypeVarTupleType):
             res.append(Constraint(tp, direction, TupleType(list(middle), tp.tuple_fallback)))
-    elif actual_unpack is not None:
-        # A special case for a variadic tuple unpack, we simply infer T <: X from
-        # Tuple[..., *tuple[T, ...], ...] <: Tuple[..., *tuple[X, ...], ...].
-        actual_unpack_type = actual_args[actual_unpack]
-        assert isinstance(actual_unpack_type, UnpackType)
-        a_unpacked = get_proper_type(actual_unpack_type.type)
-        if isinstance(a_unpacked, Instance) and a_unpacked.type.fullname == "builtins.tuple":
-            t_unpack = template_args[template_unpack]
-            assert isinstance(t_unpack, UnpackType)
-            tp = get_proper_type(t_unpack.type)
-            if isinstance(tp, Instance) and tp.type.fullname == "builtins.tuple":
-                res.extend(infer_constraints(tp.args[0], a_unpacked.args[0], direction))
     return res
 
 

mypy/subtypes.py

@@ -960,6 +960,8 @@ def variadic_tuple_subtype(self, left: TupleType, right: TupleType) -> bool:
             if not self._is_subtype(left_item, right_item):
                 return False
             max_overlap = max(0, right_prefix - left_prefix, right_suffix - left_suffix)
+            # We need to also handle the case where overlap is positive on both sides.
+            max_overlap = max(max_overlap, right_prefix - left_prefix + right_suffix - left_suffix)
             for overlap in range(max_overlap + 1):
                 repr_items = left.items[:left_prefix] + [left_item] * overlap
                 if left_suffix:

mypy/types.py

@@ -4253,7 +4253,15 @@ def has_recursive_types(typ: Type) -> bool:
 def split_with_prefix_and_suffix(
     types: tuple[Type, ...], prefix: int, suffix: int
 ) -> tuple[tuple[Type, ...], tuple[Type, ...], tuple[Type, ...]]:
-    if len(types) <= prefix + suffix:
+    # The caller must validate that the split can be satisfied, i.e. there is
+    # enough capacity in either initial type list or there is a variadic unpack.
+    # Otherwise, this function may return nonsensical result.
+    # TODO: should we add an assert here?
+    needs_extend = False
+    index = find_unpack_in_list(types)
+    if index is not None:
+        needs_extend = index < prefix or len(types) - index - 1 < suffix
+    if needs_extend:
         types = extend_args_for_prefix_and_suffix(types, prefix, suffix)
     if suffix:
         return types[:prefix], types[prefix:-suffix], types[-suffix:]

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

@@ -3104,3 +3104,114 @@ foo(t2)
 foo(t3)  # E: Argument 1 to "foo" has incompatible type "C[Unpack[tuple[Any, ...]], str]"; expected "C[int, Unpack[tuple[int, ...]], int]"
 foo(t4)
 [builtins fixtures/tuple.pyi]
+
+[case testInferAgainstVariadicAnySplit]
+from typing import Generic, TypeVar, Optional, Any
+from typing_extensions import TypeVarTuple, Unpack
+
+T = TypeVar("T", bound=Any)
+Ts = TypeVarTuple("Ts")
+
+class Result(Generic[Unpack[Ts]]):
+    def foo(self: Result[T, Unpack[tuple[Any, ...]]]) -> Optional[T]: ...
+    def bar(self: Result[Unpack[tuple[Any, ...]], T]) -> Optional[T]: ...
+
+result: Result[Unpack[tuple[Any, ...]]]
+reveal_type(result.foo)  # N: Revealed type is "def () -> Any | None"
+reveal_type(result.bar)  # N: Revealed type is "def () -> Any | None"
+[builtins fixtures/tuple.pyi]
+
+[case testInferAgainstVariadicPrefixProtocol]
+from typing import Any, Protocol, TypeVar
+from typing_extensions import TypeVarTuple, Unpack
+
+T = TypeVar("T")
+Ts = TypeVarTuple("Ts")
+
+class WithTuple(Protocol[Unpack[Ts]]):
+    tup: tuple[Unpack[Ts]]
+
+def get_tuple(e: WithTuple[Unpack[Ts]]) -> tuple[Unpack[Ts]]:
+    return e.tup
+
+def get_first(e: WithTuple[T, Unpack[tuple[Any, ...]]]) -> T:
+    return e.tup[0]
+
+class X:
+    tup: tuple[int, str]
+
+reveal_type(get_tuple(X()))  # N: Revealed type is "tuple[builtins.int, builtins.str]"
+reveal_type(get_first(X()))  # N: Revealed type is "builtins.int"
+[builtins fixtures/tuple.pyi]
+
+[case testInferAgainstVariadicPrefixSuffix]
+from typing import Any, Protocol, TypeVar
+from typing_extensions import TypeVarTuple, Unpack
+
+T = TypeVar("T")
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+T4 = TypeVar("T4")
+Ts = TypeVarTuple("Ts")
+
+def first_last(t: tuple[T1, Unpack[tuple[Any, ...]], T2]) -> tuple[T1, T2]: ...
+def first_last_2(t: tuple[T1, T2, Unpack[Ts], T3, T4]) -> tuple[tuple[T1, T2, T3, T4], tuple[Unpack[Ts]]]: ...
+def first_last_middle(t: tuple[T1, Unpack[Ts], T2]) -> tuple[tuple[T1, T2], tuple[Unpack[Ts]]]: ...
+
+t1: tuple[Any, ...]
+t2i: tuple[str, Unpack[tuple[int, ...]], bytes]
+t2a: tuple[str, Unpack[tuple[Any, ...]], bytes]
+t3: tuple[str, bytes, Unpack[tuple[int, ...]], float, bool]
+
+reveal_type(first_last(t1))  # N: Revealed type is "tuple[Any, Any]"
+reveal_type(first_last(t2i))  # N: Revealed type is "tuple[builtins.str, builtins.bytes]"
+reveal_type(first_last(t3))  # N: Revealed type is "tuple[builtins.str, builtins.bool]"
+
+reveal_type(first_last_2(t1))  # N: Revealed type is "tuple[tuple[Any, Any, Any, Any], builtins.tuple[Any, ...]]"
+reveal_type(first_last_2(t2a))  # N: Revealed type is "tuple[tuple[builtins.str, Any, Any, builtins.bytes], builtins.tuple[Any, ...]]"
+reveal_type(first_last_2(t3))  # N: Revealed type is "tuple[tuple[builtins.str, builtins.bytes, builtins.float, builtins.bool], builtins.tuple[builtins.int, ...]]"
+
+reveal_type(first_last_middle(t3))  # N: Revealed type is "tuple[tuple[builtins.str, builtins.bool], tuple[builtins.bytes, Unpack[builtins.tuple[builtins.int, ...]], builtins.float]]"
+[builtins fixtures/tuple.pyi]
+
+[case testInferAgainstVariadicPrefixSuffix2]
+from typing import Any, TypeVar, Generic
+from typing_extensions import TypeVarTuple, Unpack
+
+T = TypeVar("T")
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+Ts = TypeVarTuple("Ts")
+
+def mix(x: tuple[T1, T2, Unpack[Ts]]) -> tuple[T1, T2, tuple[Unpack[Ts]]]: ...
+a: tuple[str, Unpack[tuple[Any, ...]], bytes, float]
+
+reveal_type(mix(a))  # N: Revealed type is "tuple[builtins.str, Any, tuple[Unpack[builtins.tuple[Any, ...]], builtins.bytes, builtins.float]]"
+
+class C(Generic[T1, Unpack[Ts], T2]):
+     start: T1
+     end: T2
+     middle: tuple[Unpack[Ts]]
+
+c: C[str, bool, Unpack[tuple[int, ...]]]
+reveal_type(c)  # N: Revealed type is "__main__.C[builtins.str, builtins.bool, Unpack[builtins.tuple[builtins.int, ...]]]"
+reveal_type(c.start)  # N: Revealed type is "builtins.str"
+reveal_type(c.middle)  # N: Revealed type is "tuple[builtins.bool, Unpack[builtins.tuple[builtins.int, ...]]]"
+reveal_type(c.end)  # N: Revealed type is "builtins.int"
+
+cs: C[Unpack[tuple[int, ...]], str, bool]
+reveal_type(cs)  # N: Revealed type is "__main__.C[Unpack[builtins.tuple[builtins.int, ...]], builtins.str, builtins.bool]"
+reveal_type(cs.start)  # N: Revealed type is "builtins.int"
+reveal_type(cs.middle)  # N: Revealed type is "tuple[Unpack[builtins.tuple[builtins.int, ...]], builtins.str]"
+reveal_type(cs.end)  # N: Revealed type is "builtins.bool"
+
+def mix2(x: tuple[T1, T, T2]) -> tuple[T, tuple[T1, T2]]: ...
+b1: tuple[int, str, Unpack[tuple[Any, ...]]]
+b2: tuple[int, Unpack[tuple[Any, ...]], str]
+b3: tuple[Unpack[tuple[Any, ...]], int, str]
+
+reveal_type(mix2(b1))  # N: Revealed type is "tuple[builtins.str, tuple[builtins.int, Any]]"
+reveal_type(mix2(b2))  # N: Revealed type is "tuple[Any, tuple[builtins.int, builtins.str]]"
+reveal_type(mix2(b3))  # N: Revealed type is "tuple[builtins.int, tuple[Any, builtins.str]]"
+[builtins fixtures/tuple.pyi]