Handle prefix/suffix in typevartuple *args support (#14112)

This requires handling more cases in the various places that we
previously modified to support *args in general. We also need to refresh
the formals-to-actuals twice in checkexpr as now it can happen in the
infer_function_type_arguments_using_context call.

The handling here is kind of asymmetric, because we can convert prefices
into positional arguments, but there is no equivalent for suffices, so
we represent that as a Tuple[Unpack[...], <suffix>] and handle that case
separately in some spots.

We also support various edge cases like passing in a tuple without any
typevartuples involved.

mypy/applytype.py

@@ -4,13 +4,14 @@
 
 import mypy.subtypes
 from mypy.expandtype import expand_type, expand_unpack_with_variables
-from mypy.nodes import ARG_POS, ARG_STAR, Context
+from mypy.nodes import ARG_STAR, Context
 from mypy.types import (
  AnyType,
  CallableType,
  Parameters,
  ParamSpecType,
  PartialType,
+ TupleType,
  Type,
  TypeVarId,
  TypeVarLikeType,
@@ -19,6 +20,7 @@
  UnpackType,
  get_proper_type,
 )
+from mypy.typevartuples import find_unpack_in_list, replace_starargs
 
 
 def get_target_type(
@@ -114,39 +116,57 @@ 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):
- expanded = expand_unpack_with_variables(var_arg.typ, id_to_type)
- assert isinstance(expanded, list)
- # Handle other cases later.
- for t in expanded:
- assert not isinstance(t, UnpackType)
  star_index = callable.arg_kinds.index(ARG_STAR)
- arg_kinds = (
- callable.arg_kinds[:star_index]
- + [ARG_POS] * len(expanded)
- + callable.arg_kinds[star_index + 1 :]
+ callable = callable.copy_modified(
+ arg_types=(
+ [
+ expand_type(at, id_to_type, allow_erased_callables)
+ for at in callable.arg_types[:star_index]
+ ]
+ + [callable.arg_types[star_index]]
+ + [
+ expand_type(at, id_to_type, allow_erased_callables)
+ for at in callable.arg_types[star_index + 1 :]
+ ]
+ )
  )
- arg_names = (
- callable.arg_names[:star_index]
- + [None] * len(expanded)
- + callable.arg_names[star_index + 1 :]
- )
- arg_types = (
- [
- expand_type(at, id_to_type, allow_erased_callables)
- for at in callable.arg_types[:star_index]
- ]
- + expanded
- + [
- expand_type(at, id_to_type, allow_erased_callables)
- 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)
+ assert isinstance(expanded_tvt, list)
+ for t in expanded_tvt:
+ assert not isinstance(t, UnpackType)
+ callable = replace_starargs(callable, expanded_tvt)
+ else:
+ assert False, "mypy bug: unhandled case applying unpack"
+ else:
+ callable = callable.copy_modified(
+ arg_types=[
+ expand_type(at, id_to_type, allow_erased_callables) for at in callable.arg_types
  ]
  )
- else:
- arg_types = [
- expand_type(at, id_to_type, allow_erased_callables) for at in callable.arg_types
- ]
- arg_kinds = callable.arg_kinds
- arg_names = callable.arg_names
 
  # Apply arguments to TypeGuard if any.
  if callable.type_guard is not None:
@@ -158,10 +178,7 @@ def apply_generic_arguments(
  remaining_tvars = [tv for tv in tvars if tv.id not in id_to_type]
 
  return callable.copy_modified(
- arg_types=arg_types,
  ret_type=expand_type(callable.ret_type, id_to_type, allow_erased_callables),
  variables=remaining_tvars,
  type_guard=type_guard,
- arg_kinds=arg_kinds,
- arg_names=arg_names,
  )

mypy/checker.py

@@ -1178,12 +1178,17 @@ def check_func_def(
  if isinstance(arg_type, ParamSpecType):
  pass
  elif isinstance(arg_type, UnpackType):
- arg_type = TupleType(
- [arg_type],
- fallback=self.named_generic_type(
- "builtins.tuple", [self.named_type("builtins.object")]
- ),
- )
+ if isinstance(get_proper_type(arg_type.type), TupleType):
+ # Instead of using Tuple[Unpack[Tuple[...]]], just use
+ # Tuple[...]
+ arg_type = arg_type.type
+ else:
+ arg_type = TupleType(
+ [arg_type],
+ fallback=self.named_generic_type(
+ "builtins.tuple", [self.named_type("builtins.object")]
+ ),
+ )
  else:
  # builtins.tuple[T] is typing.Tuple[T, ...]
  arg_type = self.named_generic_type("builtins.tuple", [arg_type])

mypy/checkexpr.py

@@ -150,6 +150,7 @@
  TypeVarType,
  UninhabitedType,
  UnionType,
+ UnpackType,
  flatten_nested_unions,
  get_proper_type,
  get_proper_types,
@@ -1404,13 +1405,21 @@ def check_callable_call(
  )
  callee = freshen_function_type_vars(callee)
  callee = self.infer_function_type_arguments_using_context(callee, context)
+ if need_refresh:
+ # Argument kinds etc. may have changed due to
+ # ParamSpec or TypeVarTuple variables being replaced with an arbitrary
+ # number of arguments; recalculate actual-to-formal map
+ formal_to_actual = map_actuals_to_formals(
+ arg_kinds,
+ arg_names,
+ callee.arg_kinds,
+ callee.arg_names,
+ lambda i: self.accept(args[i]),
+ )
  callee = self.infer_function_type_arguments(
  callee, args, arg_kinds, formal_to_actual, context
  )
  if need_refresh:
- # Argument kinds etc. may have changed due to
- # ParamSpec variables being replaced with an arbitrary
- # number of arguments; recalculate actual-to-formal map
  formal_to_actual = map_actuals_to_formals(
  arg_kinds,
  arg_names,
@@ -1999,11 +2008,66 @@ def check_argument_types(
  # Keep track of consumed tuple *arg items.
  mapper = ArgTypeExpander(self.argument_infer_context())
  for i, actuals in enumerate(formal_to_actual):
- for actual in actuals:
- actual_type = arg_types[actual]
+ orig_callee_arg_type = get_proper_type(callee.arg_types[i])
+
+ # Checking the case that we have more than one item but the first argument
+ # is an unpack, so this would be something like:
+ # [Tuple[Unpack[Ts]], int]
+ #
+ # In this case we have to check everything together, we do this by re-unifying
+ # the suffices to the tuple, e.g. a single actual like
+ # Tuple[Unpack[Ts], int]
+ expanded_tuple = False
+ if len(actuals) > 1:
+ first_actual_arg_type = get_proper_type(arg_types[actuals[0]])
+ if (
+ isinstance(first_actual_arg_type, TupleType)
+ and len(first_actual_arg_type.items) == 1
+ and isinstance(get_proper_type(first_actual_arg_type.items[0]), UnpackType)
+ ):
+ # TODO: use walrus operator
+ actual_types = [first_actual_arg_type.items[0]] + [
+ arg_types[a] for a in actuals[1:]
+ ]
+ 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
+ callee_arg_kinds = [nodes.ARG_STAR] + [nodes.ARG_POS] * (
+ len(orig_callee_arg_type.items) - 1
+ )
+ expanded_tuple = True
+
+ if not expanded_tuple:
+ actual_types = [arg_types[a] for a in actuals]
+ actual_kinds = [arg_kinds[a] for a in actuals]
+ if isinstance(orig_callee_arg_type, UnpackType):
+ unpacked_type = get_proper_type(orig_callee_arg_type.type)
+ # Only case we know of thus far.
+ assert isinstance(unpacked_type, TupleType)
+ actual_types = [arg_types[a] for a in actuals]
+ actual_kinds = [arg_kinds[a] for a in actuals]
+ callee_arg_types = unpacked_type.items
+ callee_arg_kinds = [ARG_POS] * len(actuals)
+ else:
+ callee_arg_types = [orig_callee_arg_type] * len(actuals)
+ callee_arg_kinds = [callee.arg_kinds[i]] * len(actuals)
+
+ assert len(actual_types) == len(actuals) == len(actual_kinds)
+
+ if len(callee_arg_types) != len(actual_types):
+ # TODO: Improve error message
+ self.chk.fail("Invalid number of arguments", context)
+ continue
+
+ assert len(callee_arg_types) == len(actual_types)
+ assert len(callee_arg_types) == len(callee_arg_kinds)
+ for actual, actual_type, actual_kind, callee_arg_type, callee_arg_kind in zip(
+ actuals, actual_types, actual_kinds, callee_arg_types, callee_arg_kinds
+ ):
  if actual_type is None:
  continue # Some kind of error was already reported.
- actual_kind = arg_kinds[actual]
  # Check that a *arg is valid as varargs.
  if actual_kind == nodes.ARG_STAR and not self.is_valid_var_arg(actual_type):
  self.msg.invalid_var_arg(actual_type, context)
@@ -2013,13 +2077,13 @@ def check_argument_types(
  is_mapping = is_subtype(actual_type, self.chk.named_type("typing.Mapping"))
  self.msg.invalid_keyword_var_arg(actual_type, is_mapping, context)
  expanded_actual = mapper.expand_actual_type(
- actual_type, actual_kind, callee.arg_names[i], callee.arg_kinds[i]
+ actual_type, actual_kind, callee.arg_names[i], callee_arg_kind
  )
  check_arg(
  expanded_actual,
  actual_type,
- arg_kinds[actual],
- callee.arg_types[i],
+ actual_kind,
+ callee_arg_type,
  actual + 1,
  i + 1,
  callee,
@@ -4719,6 +4783,7 @@ def is_valid_var_arg(self, typ: Type) -> bool:
  )
  or isinstance(typ, AnyType)
  or isinstance(typ, ParamSpecType)
+ or isinstance(typ, UnpackType)
  )
 
  def is_valid_keyword_var_arg(self, typ: Type) -> bool:

mypy/constraints.py

@@ -133,8 +133,26 @@ def infer_constraints_for_callable(
  )
  )
 
- assert isinstance(unpack_type.type, TypeVarTupleType)
- constraints.append(Constraint(unpack_type.type, SUPERTYPE_OF, TypeList(actual_types)))
+ unpacked_type = get_proper_type(unpack_type.type)
+ if isinstance(unpacked_type, TypeVarTupleType):
+ constraints.append(Constraint(unpacked_type, SUPERTYPE_OF, TypeList(actual_types)))
+ elif isinstance(unpacked_type, TupleType):
+ # Prefixes get converted to positional args, so technically the only case we
+ # should have here is like Tuple[Unpack[Ts], Y1, Y2, Y3]. If this turns out
+ # not to hold we can always handle the prefixes too.
+ inner_unpack = unpacked_type.items[0]
+ assert isinstance(inner_unpack, UnpackType)
+ inner_unpacked_type = get_proper_type(inner_unpack.type)
+ assert isinstance(inner_unpacked_type, TypeVarTupleType)
+ suffix_len = len(unpacked_type.items) - 1
+ constraints.append(
+ Constraint(
+ inner_unpacked_type, SUPERTYPE_OF, TypeList(actual_types[:-suffix_len])
+ )
+ )
+ else:
+ assert False, "mypy bug: unhandled constraint inference case"
+
  else:
  for actual in actuals:
  actual_arg_type = arg_types[actual]