Support user defined variadic tuple types (#15961)

Fixes #15946

Note this actually adds support also for variadic NamedTuples and
variadic TypedDicts. Not that anyone requested this, but since generic
NamedTuples and generic TypedDicts are supported using the same
mechanism (special aliases) as generic tuple types (like `class
A(Tuple[T, S]): ...` in the issue), it looked more risky and arbitrary
to _not_support them.

Btw the implementation is simple, but while I was working on this, I
accidentally found a problem with my general idea of doing certain type
normlaizations in `semanal_typeargs.py`. The problem is that sometimes
we can call `get_proper_type()` during semantic analysis, so all the
code that gets triggered by this (mostly `expand_type()`) can't really
rely on types being normalized. Fortunately, with just few tweaks I
manged to make the code mostly robust to such scenarios (TBH there are
few possible holes left, but this is getting really complex, I think it
is better to release this, and see if people will ever hit such
scenarios, then fix accordingly).

Author: ilevkivskyi

mypy/expandtype.py

@@ -269,7 +269,8 @@ def visit_unpack_type(self, t: UnpackType) -> Type:
         # instead.
         # However, if the item is a variadic tuple, we can simply carry it over.
         # In particular, if we expand A[*tuple[T, ...]] with substitutions {T: str},
-        # it is hard to assert this without getting proper type.
+        # it is hard to assert this without getting proper type. Another important
+        # example is non-normalized types when called from semanal.py.
         return UnpackType(t.type.accept(self))
 
     def expand_unpack(self, t: UnpackType) -> list[Type] | AnyType | UninhabitedType:
@@ -414,6 +415,10 @@ def visit_tuple_type(self, t: TupleType) -> Type:
                     unpacked = get_proper_type(item.type)
                     if isinstance(unpacked, Instance):
                         assert unpacked.type.fullname == "builtins.tuple"
+                        if t.partial_fallback.type.fullname != "builtins.tuple":
+                            # If it is a subtype (like named tuple) we need to preserve it,
+                            # this essentially mimics the logic in tuple_fallback().
+                            return t.partial_fallback.accept(self)
                         return unpacked
             fallback = t.partial_fallback.accept(self)
             assert isinstance(fallback, ProperType) and isinstance(fallback, Instance)

mypy/maptype.py

@@ -113,6 +113,5 @@ def instance_to_type_environment(instance: Instance) -> dict[TypeVarId, Type]:
     required number of type arguments.  So this environment consists
     of the class's type variables mapped to the Instance's actual
     arguments.  The type variables are mapped by their `id`.
-
     """
     return {binder.id: arg for binder, arg in zip(instance.type.defn.type_vars, instance.args)}

mypy/nodes.py

@@ -3546,7 +3546,12 @@ def from_tuple_type(cls, info: TypeInfo) -> TypeAlias:
         assert info.tuple_type
         # TODO: is it possible to refactor this to set the correct type vars here?
         return TypeAlias(
-            info.tuple_type.copy_modified(fallback=mypy.types.Instance(info, info.defn.type_vars)),
+            info.tuple_type.copy_modified(
+                # Create an Instance similar to fill_typevars().
+                fallback=mypy.types.Instance(
+                    info, mypy.types.type_vars_as_args(info.defn.type_vars)
+                )
+            ),
             info.fullname,
             info.line,
             info.column,
@@ -3563,7 +3568,10 @@ def from_typeddict_type(cls, info: TypeInfo) -> TypeAlias:
         # TODO: is it possible to refactor this to set the correct type vars here?
         return TypeAlias(
             info.typeddict_type.copy_modified(
-                fallback=mypy.types.Instance(info, info.defn.type_vars)
+                # Create an Instance similar to fill_typevars().
+                fallback=mypy.types.Instance(
+                    info, mypy.types.type_vars_as_args(info.defn.type_vars)
+                )
             ),
             info.fullname,
             info.line,

mypy/semanal.py

@@ -277,6 +277,7 @@
     get_proper_types,
     is_named_instance,
     remove_dups,
+    type_vars_as_args,
 )
 from mypy.types_utils import is_invalid_recursive_alias, store_argument_type
 from mypy.typevars import fill_typevars
@@ -1702,12 +1703,17 @@ def setup_type_vars(self, defn: ClassDef, tvar_defs: list[TypeVarLikeType]) -> N
     def setup_alias_type_vars(self, defn: ClassDef) -> None:
         assert defn.info.special_alias is not None
         defn.info.special_alias.alias_tvars = list(defn.type_vars)
+        # It is a bit unfortunate that we need to inline some logic from TypeAlias constructor,
+        # but it is required, since type variables may change during semantic analyzer passes.
+        for i, t in enumerate(defn.type_vars):
+            if isinstance(t, TypeVarTupleType):
+                defn.info.special_alias.tvar_tuple_index = i
         target = defn.info.special_alias.target
         assert isinstance(target, ProperType)
         if isinstance(target, TypedDictType):
-            target.fallback.args = tuple(defn.type_vars)
+            target.fallback.args = type_vars_as_args(defn.type_vars)
         elif isinstance(target, TupleType):
-            target.partial_fallback.args = tuple(defn.type_vars)
+            target.partial_fallback.args = type_vars_as_args(defn.type_vars)
         else:
             assert False, f"Unexpected special alias type: {type(target)}"
 

mypy/semanal_typeargs.py

@@ -86,31 +86,31 @@ def visit_type_alias_type(self, t: TypeAliasType) -> None:
         # correct aliases. Also, variadic aliases are better to check when fully analyzed,
         # so we do this here.
         assert t.alias is not None, f"Unfixed type alias {t.type_ref}"
-        args = flatten_nested_tuples(t.args)
+        # TODO: consider moving this validation to typeanal.py, expanding invalid aliases
+        # during semantic analysis may cause crashes.
         if t.alias.tvar_tuple_index is not None:
-            correct = len(args) >= len(t.alias.alias_tvars) - 1
+            correct = len(t.args) >= len(t.alias.alias_tvars) - 1
             if any(
                 isinstance(a, UnpackType) and isinstance(get_proper_type(a.type), Instance)
-                for a in args
+                for a in t.args
             ):
                 correct = True
         else:
-            correct = len(args) == len(t.alias.alias_tvars)
+            correct = len(t.args) == len(t.alias.alias_tvars)
         if not correct:
             if t.alias.tvar_tuple_index is not None:
                 exp_len = f"at least {len(t.alias.alias_tvars) - 1}"
             else:
                 exp_len = f"{len(t.alias.alias_tvars)}"
             self.fail(
-                f"Bad number of arguments for type alias, expected: {exp_len}, given: {len(args)}",
+                "Bad number of arguments for type alias,"
+                f" expected: {exp_len}, given: {len(t.args)}",
                 t,
                 code=codes.TYPE_ARG,
             )
             t.args = set_any_tvars(
                 t.alias, t.line, t.column, self.options, from_error=True, fail=self.fail
             ).args
-        else:
-            t.args = args
         is_error = self.validate_args(t.alias.name, t.args, t.alias.alias_tvars, t)
         if not is_error:
             # If there was already an error for the alias itself, there is no point in checking

mypy/semanal_typeddict.py

@@ -252,6 +252,7 @@ def map_items_to_base(
             if not tvars:
                 mapped_items[key] = type_in_base
                 continue
+            # TODO: simple zip can't be used for variadic types.
             mapped_items[key] = expand_type(
                 type_in_base, {t.id: a for (t, a) in zip(tvars, base_args)}
             )

mypy/typeanal.py

@@ -82,6 +82,7 @@
     UnionType,
     UnpackType,
     callable_with_ellipsis,
+    flatten_nested_tuples,
     flatten_nested_unions,
     get_proper_type,
     has_type_vars,
@@ -763,8 +764,8 @@ def analyze_type_with_type_info(
             if info.special_alias:
                 return instantiate_type_alias(
                     info.special_alias,
-                    # TODO: should we allow NamedTuples generic in ParamSpec and TypeVarTuple?
-                    self.anal_array(args),
+                    # TODO: should we allow NamedTuples generic in ParamSpec?
+                    self.anal_array(args, allow_unpack=True),
                     self.fail,
                     False,
                     ctx,
@@ -782,7 +783,7 @@ def analyze_type_with_type_info(
                 return instantiate_type_alias(
                     info.special_alias,
                     # TODO: should we allow TypedDicts generic in ParamSpec?
-                    self.anal_array(args),
+                    self.anal_array(args, allow_unpack=True),
                     self.fail,
                     False,
                     ctx,
@@ -1948,7 +1949,10 @@ def instantiate_type_alias(
     # TODO: we need to check args validity w.r.t alias.alias_tvars.
     # Otherwise invalid instantiations will be allowed in runtime context.
     # Note: in type context, these will be still caught by semanal_typeargs.
-    typ = TypeAliasType(node, args, ctx.line, ctx.column)
+    # Type aliases are special, since they can be expanded during semantic analysis,
+    # so we need to normalize them as soon as possible.
+    # TODO: can this cause an infinite recursion?
+    typ = TypeAliasType(node, flatten_nested_tuples(args), ctx.line, ctx.column)
     assert typ.alias is not None
     # HACK: Implement FlexibleAlias[T, typ] by expanding it to typ here.
     if (

mypy/types.py

@@ -1045,9 +1045,12 @@ class UnpackType(ProperType):
     """Type operator Unpack from PEP646. Can be either with Unpack[]
     or unpacking * syntax.
 
-    The inner type should be either a TypeVarTuple, a constant size
-    tuple, or a variable length tuple. Type aliases to these are not allowed,
-    except during semantic analysis.
+    The inner type should be either a TypeVarTuple, or a variable length tuple.
+    In an exceptional case of callable star argument it can be a fixed length tuple.
+
+    Note: the above restrictions are only guaranteed by normalizations after semantic
+    analysis, if your code needs to handle UnpackType *during* semantic analysis, it is
+    wild west, technically anything can be present in the wrapped type.
     """
 
     __slots__ = ["type"]
@@ -2143,7 +2146,11 @@ def with_normalized_var_args(self) -> Self:
                     assert nested_unpacked.type.fullname == "builtins.tuple"
                     new_unpack = nested_unpacked.args[0]
                 else:
-                    assert isinstance(nested_unpacked, TypeVarTupleType)
+                    if not isinstance(nested_unpacked, TypeVarTupleType):
+                        # We found a non-nomralized tuple type, this means this method
+                        # is called during semantic analysis (e.g. from get_proper_type())
+                        # there is no point in normalizing callables at this stage.
+                        return self
                     new_unpack = nested_unpack
             else:
                 new_unpack = UnpackType(
@@ -3587,6 +3594,17 @@ def remove_dups(types: list[T]) -> list[T]:
     return new_types
 
 
+def type_vars_as_args(type_vars: Sequence[TypeVarLikeType]) -> tuple[Type, ...]:
+    """Represent type variables as they would appear in a type argument list."""
+    args: list[Type] = []
+    for tv in type_vars:
+        if isinstance(tv, TypeVarTupleType):
+            args.append(UnpackType(tv))
+        else:
+            args.append(tv)
+    return tuple(args)
+
+
 # This cyclic import is unfortunate, but to avoid it we would need to move away all uses
 # of get_proper_type() from types.py. Majority of them have been removed, but few remaining
 # are quite tricky to get rid of, but ultimately we want to do it at some point.

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

@@ -1032,3 +1032,90 @@ Second = Tuple[C, D]
 x: G[Unpack[First], Unpack[Second]]  # E: Type argument "A" of "G" must be a subtype of "int" \
                                      # E: Type argument "D" of "G" must be a subtype of "str"
 [builtins fixtures/tuple.pyi]
+
+[case testVariadicTupleType]
+from typing import Tuple, Callable
+from typing_extensions import TypeVarTuple, Unpack
+
+Ts = TypeVarTuple("Ts")
+class A(Tuple[Unpack[Ts]]):
+    fn: Callable[[Unpack[Ts]], None]
+
+x: A[int]
+reveal_type(x)  # N: Revealed type is "Tuple[builtins.int, fallback=__main__.A[builtins.int]]"
+reveal_type(x[0])  # N: Revealed type is "builtins.int"
+reveal_type(x.fn)  # N: Revealed type is "def (builtins.int)"
+
+y: A[int, str]
+reveal_type(y)  # N: Revealed type is "Tuple[builtins.int, builtins.str, fallback=__main__.A[builtins.int, builtins.str]]"
+reveal_type(y[0])  # N: Revealed type is "builtins.int"
+reveal_type(y.fn)  # N: Revealed type is "def (builtins.int, builtins.str)"
+
+z: A[Unpack[Tuple[int, ...]]]
+reveal_type(z)  # N: Revealed type is "__main__.A[Unpack[builtins.tuple[builtins.int, ...]]]"
+# TODO: this requires fixing map_instance_to_supertype().
+# reveal_type(z[0])
+reveal_type(z.fn)  # N: Revealed type is "def (*builtins.int)"
+
+t: A[int, Unpack[Tuple[int, str]], str]
+reveal_type(t)  # N: Revealed type is "Tuple[builtins.int, builtins.int, builtins.str, builtins.str, fallback=__main__.A[builtins.int, builtins.int, builtins.str, builtins.str]]"
+reveal_type(t[0])  # N: Revealed type is "builtins.int"
+reveal_type(t.fn)  # N: Revealed type is "def (builtins.int, builtins.int, builtins.str, builtins.str)"
+[builtins fixtures/tuple.pyi]
+
+[case testVariadicNamedTuple]
+from typing import Tuple, Callable, NamedTuple, Generic
+from typing_extensions import TypeVarTuple, Unpack
+
+Ts = TypeVarTuple("Ts")
+class A(NamedTuple, Generic[Unpack[Ts], T]):
+    fn: Callable[[Unpack[Ts]], None]
+    val: T
+
+y: A[int, str]
+reveal_type(y)  # N: Revealed type is "Tuple[def (builtins.int), builtins.str, fallback=__main__.A[builtins.int, builtins.str]]"
+reveal_type(y[0])  # N: Revealed type is "def (builtins.int)"
+reveal_type(y.fn)  # N: Revealed type is "def (builtins.int)"
+
+z: A[Unpack[Tuple[int, ...]]]
+reveal_type(z)  # N: Revealed type is "Tuple[def (*builtins.int), builtins.int, fallback=__main__.A[Unpack[builtins.tuple[builtins.int, ...]], builtins.int]]"
+reveal_type(z.fn)  # N: Revealed type is "def (*builtins.int)"
+
+t: A[int, Unpack[Tuple[int, str]], str]
+reveal_type(t)  # N: Revealed type is "Tuple[def (builtins.int, builtins.int, builtins.str), builtins.str, fallback=__main__.A[builtins.int, builtins.int, builtins.str, builtins.str]]"
+
+def test(x: int, y: str) -> None: ...
+nt = A(fn=test, val=42)
+reveal_type(nt)  # N: Revealed type is "Tuple[def (builtins.int, builtins.str), builtins.int, fallback=__main__.A[builtins.int, builtins.str, builtins.int]]"
+
+def bad() -> int: ...
+nt2 = A(fn=bad, val=42)  # E: Argument "fn" to "A" has incompatible type "Callable[[], int]"; expected "Callable[[], None]"
+[builtins fixtures/tuple.pyi]
+
+[case testVariadicTypedDict]
+from typing import Tuple, Callable, Generic
+from typing_extensions import TypeVarTuple, Unpack, TypedDict
+
+Ts = TypeVarTuple("Ts")
+class A(TypedDict, Generic[Unpack[Ts], T]):
+    fn: Callable[[Unpack[Ts]], None]
+    val: T
+
+y: A[int, str]
+reveal_type(y)  # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int), 'val': builtins.str})"
+reveal_type(y["fn"])  # N: Revealed type is "def (builtins.int)"
+
+z: A[Unpack[Tuple[int, ...]]]
+reveal_type(z)  # N: Revealed type is "TypedDict('__main__.A', {'fn': def (*builtins.int), 'val': builtins.int})"
+reveal_type(z["fn"])  # N: Revealed type is "def (*builtins.int)"
+
+t: A[int, Unpack[Tuple[int, str]], str]
+reveal_type(t)  # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int, builtins.int, builtins.str), 'val': builtins.str})"
+
+def test(x: int, y: str) -> None: ...
+td = A({"fn": test, "val": 42})
+reveal_type(td)  # N: Revealed type is "TypedDict('__main__.A', {'fn': def (builtins.int, builtins.str), 'val': builtins.int})"
+
+def bad() -> int: ...
+td2 = A({"fn": bad, "val": 42})  # E: Incompatible types (expression has type "Callable[[], int]", TypedDict item "fn" has type "Callable[[], None]")
+[builtins fixtures/tuple.pyi]