Polymorphic inference: basic support for variadic types

mypy/applytype.py

@@ -3,15 +3,13 @@
 from typing import Callable, Sequence
 
 import mypy.subtypes
-from mypy.expandtype import expand_type, expand_unpack_with_variables
-from mypy.nodes import ARG_STAR, Context
+from mypy.expandtype import expand_type
+from mypy.nodes import Context
 from mypy.types import (
     AnyType,
     CallableType,
-    Instance,
     ParamSpecType,
     PartialType,
-    TupleType,
     Type,
     TypeVarId,
     TypeVarLikeType,
@@ -21,7 +19,6 @@
     UnpackType,
     get_proper_type,
 )
-from mypy.typevartuples import find_unpack_in_list, replace_starargs
 
 
 def get_target_type(
@@ -107,6 +104,8 @@ def apply_generic_arguments(
         if target_type is not None:
             id_to_type[tvar.id] = target_type
 
+    # TODO: validate arg_kinds/arg_names for ParamSpec and TypeVarTuple replacements,
+    # not just type variable bounds above.
     param_spec = callable.param_spec()
     if param_spec is not None:
         nt = id_to_type.get(param_spec.id)
@@ -122,55 +121,9 @@ def apply_generic_arguments(
     # Apply arguments to argument types.
     var_arg = callable.var_arg()
     if var_arg is not None and isinstance(var_arg.typ, UnpackType):
-        star_index = callable.arg_kinds.index(ARG_STAR)
-        callable = callable.copy_modified(
-            arg_types=(
-                [expand_type(at, id_to_type) for at in callable.arg_types[:star_index]]
-                + [callable.arg_types[star_index]]
-                + [expand_type(at, id_to_type) for at in callable.arg_types[star_index + 1 :]]
-            )
-        )
-
-        unpacked_type = get_proper_type(var_arg.typ.type)
-        if isinstance(unpacked_type, TupleType):
-            # Assuming for now that because we convert prefixes to positional arguments,
-            # the first argument is always an unpack.
-            expanded_tuple = expand_type(unpacked_type, id_to_type)
-            if isinstance(expanded_tuple, TupleType):
-                # TODO: handle the case where the tuple has an unpack. This will
-                # hit an assert below.
-                expanded_unpack = find_unpack_in_list(expanded_tuple.items)
-                if expanded_unpack is not None:
-                    callable = callable.copy_modified(
-                        arg_types=(
-                            callable.arg_types[:star_index]
-                            + [expanded_tuple]
-                            + callable.arg_types[star_index + 1 :]
-                        )
-                    )
-                else:
-                    callable = replace_starargs(callable, expanded_tuple.items)
-            else:
-                # TODO: handle the case for if we get a variable length tuple.
-                assert False, f"mypy bug: unimplemented case, {expanded_tuple}"
-        elif isinstance(unpacked_type, TypeVarTupleType):
-            expanded_tvt = expand_unpack_with_variables(var_arg.typ, id_to_type)
-            if isinstance(expanded_tvt, list):
-                for t in expanded_tvt:
-                    assert not isinstance(t, UnpackType)
-                callable = replace_starargs(callable, expanded_tvt)
-            else:
-                assert isinstance(expanded_tvt, Instance)
-                assert expanded_tvt.type.fullname == "builtins.tuple"
-                callable = callable.copy_modified(
-                    arg_types=(
-                        callable.arg_types[:star_index]
-                        + [expanded_tvt.args[0]]
-                        + callable.arg_types[star_index + 1 :]
-                    )
-                )
-        else:
-            assert False, "mypy bug: unhandled case applying unpack"
+        callable = expand_type(callable, id_to_type)
+        assert isinstance(callable, CallableType)
+        return callable.copy_modified(variables=[tv for tv in tvars if tv.id not in id_to_type])
     else:
         callable = callable.copy_modified(
             arg_types=[expand_type(at, id_to_type) for at in callable.arg_types]
@@ -183,6 +136,9 @@ def apply_generic_arguments(
         type_guard = None
 
     # The callable may retain some type vars if only some were applied.
+    # TODO: move apply_poly() logic from checkexpr.py here when new inference
+    # becomes universally used (i.e. in all passes + in unification).
+    # With this new logic we can actually *add* some new free variables.
     remaining_tvars = [tv for tv in tvars if tv.id not in id_to_type]
 
     return callable.copy_modified(

mypy/checkexpr.py

@@ -2312,11 +2312,15 @@ def check_argument_types(
                     ]
                     actual_kinds = [nodes.ARG_STAR] + [nodes.ARG_POS] * (len(actuals) - 1)
 
-                    assert isinstance(orig_callee_arg_type, TupleType)
-                    assert orig_callee_arg_type.items
-                    callee_arg_types = orig_callee_arg_type.items
+                    # TODO: can we really assert this? What if formal is just plain Unpack[Ts]?
+                    assert isinstance(orig_callee_arg_type, UnpackType)
+                    assert isinstance(orig_callee_arg_type.type, ProperType) and isinstance(
+                        orig_callee_arg_type.type, TupleType
+                    )
+                    assert orig_callee_arg_type.type.items
+                    callee_arg_types = orig_callee_arg_type.type.items
                     callee_arg_kinds = [nodes.ARG_STAR] + [nodes.ARG_POS] * (
-                        len(orig_callee_arg_type.items) - 1
+                        len(orig_callee_arg_type.type.items) - 1
                     )
                     expanded_tuple = True
 
@@ -5792,8 +5796,9 @@ def visit_param_spec(self, t: ParamSpecType) -> Type:
         return super().visit_param_spec(t)
 
     def visit_type_var_tuple(self, t: TypeVarTupleType) -> Type:
-        # TODO: Support polymorphic apply for TypeVarTuple.
-        raise PolyTranslationError()
+        if t in self.poly_tvars and t not in self.bound_tvars:
+            raise PolyTranslationError()
+        return super().visit_type_var_tuple(t)
 
     def visit_type_alias_type(self, t: TypeAliasType) -> Type:
         if not t.args:
@@ -5827,7 +5832,6 @@ def visit_instance(self, t: Instance) -> Type:
                 return t.copy_modified(args=new_args)
         # There is the same problem with callback protocols as with aliases
         # (callback protocols are essentially more flexible aliases to callables).
-        # Note: consider supporting bindings in instances, e.g. LRUCache[[x: T], T].
         if t.args and t.type.is_protocol and t.type.protocol_members == ["__call__"]:
             if t.type in self.seen_aliases:
                 raise PolyTranslationError()
@@ -5862,6 +5866,12 @@ def __init__(self) -> None:
     def visit_type_var(self, t: TypeVarType) -> bool:
         return True
 
+    def visit_param_spec(self, t: ParamSpecType) -> bool:
+        return True
+
+    def visit_type_var_tuple(self, t: TypeVarTupleType) -> bool:
+        return True
+
 
 def has_erased_component(t: Type | None) -> bool:
     return t is not None and t.accept(HasErasedComponentsQuery())