diff --git a/crates/ty_python_semantic/resources/mdtest/annotations/starred.md b/crates/ty_python_semantic/resources/mdtest/annotations/starred.md index 2d5e4dd108a330..edf5145dfff59e 100644 --- a/crates/ty_python_semantic/resources/mdtest/annotations/starred.md +++ b/crates/ty_python_semantic/resources/mdtest/annotations/starred.md @@ -17,10 +17,8 @@ def append_int(*args: *Ts) -> tuple[*Ts, int]: return (*args, 1) -# TODO should be tuple[Literal[True], Literal["a"], int] -reveal_type(append_int(True, "a")) # revealed: tuple[*tuple[Unknown, ...], int] -# TODO should be tuple[int] -reveal_type(append_int()) # revealed: tuple[*tuple[Unknown, ...], int] +reveal_type(append_int(True, "a")) # revealed: tuple[Literal[True], Literal["a"], int] +reveal_type(append_int()) # revealed: tuple[int] def first_arg_int(*args: *tuple[int, *tuple[str, ...]]): ... diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/typevartuple.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/typevartuple.md index 832cb3dda7ba9b..e8ef88221624b6 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/typevartuple.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/typevartuple.md @@ -370,13 +370,12 @@ class Variadic(Generic[*Ts]): reveal_type(Positional(())) # revealed: Positional[()] reveal_type(Positional((1, "a"))) # revealed: Positional[int, str] -# TODO: Infer the `TypeVarTuple` from arguments matched to the variadic parameter. -reveal_type(Variadic()) # revealed: Variadic[*tuple[Unknown, ...]] -reveal_type(Variadic(1, "a")) # revealed: Variadic[*tuple[Unknown, ...]] +reveal_type(Variadic()) # revealed: Variadic[()] +reveal_type(Variadic(1, "a")) # revealed: Variadic[int, str] def _(i: int, s: str) -> None: reveal_type(Positional((i, s))) # revealed: Positional[int, str] - reveal_type(Variadic(i, s)) # revealed: Variadic[*tuple[Unknown, ...]] + reveal_type(Variadic(i, s)) # revealed: Variadic[int, str] ``` ### Unspecified type arguments @@ -445,6 +444,46 @@ class WithBackportedDefault(Generic[Unpack[Ts]]): reveal_type(WithBackportedDefault().attr) # revealed: tuple[int, str] ``` +## Generic Functions + +### Starred variadic parameters + +A legacy type-variable tuple is inferred from all positional arguments matched to `*args`, including +an empty argument list. + +```py +from typing import TypeVarTuple, assert_type + +Ts = TypeVarTuple("Ts") + +def args_to_tuple(*args: *Ts) -> tuple[*Ts]: + raise NotImplementedError + +def _(i: int, s: str) -> None: + assert_type(args_to_tuple(), tuple[()]) + assert_type(args_to_tuple(i, s), tuple[int, str]) +``` + +### Starred variadic parameters with fixed suffixes + +Required tuple elements following the type-variable tuple are excluded from its inferred +specialization. The type-variable tuple can still be empty. + +```py +from typing import TypeVarTuple, assert_type + +Ts = TypeVarTuple("Ts") + +class Env: ... + +def exec_le(path: str, *args: *tuple[*Ts, Env], env: Env | None = None) -> tuple[*Ts]: + raise NotImplementedError + +def _(i: int, s: str) -> None: + assert_type(exec_le("", Env()), tuple[()]) + assert_type(exec_le(s, i, s, Env()), tuple[int, str]) +``` + ## Type Aliases ### Legacy generic aliases diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/unpack.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/unpack.md index 699e4a77d06fe6..12c38697872716 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/unpack.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/unpack.md @@ -40,9 +40,8 @@ def collect(*args: Unpack[Ts]) -> tuple[Unpack[Ts]]: reveal_type(args) # revealed: tuple[*Ts@collect] raise NotImplementedError -# TODO: Infer the `TypeVarTuple` from arguments matched to the variadic parameter. -reveal_type(collect()) # revealed: tuple[Unknown, ...] -reveal_type(collect(1, "a")) # revealed: tuple[Unknown, ...] +reveal_type(collect()) # revealed: tuple[()] +reveal_type(collect(1, "a")) # revealed: tuple[Literal[1], Literal["a"]] ``` ## Callable parameters diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/typevartuple.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/typevartuple.md index dea784f91f1427..e4fc73c196c741 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/typevartuple.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/typevartuple.md @@ -135,13 +135,12 @@ class Variadic[*Ts]: reveal_type(Positional(())) # revealed: Positional[()] reveal_type(Positional((1, "a"))) # revealed: Positional[int, str] -# TODO: Infer the `TypeVarTuple` from arguments matched to the variadic parameter. -reveal_type(Variadic()) # revealed: Variadic[*tuple[Unknown, ...]] -reveal_type(Variadic(1, "a")) # revealed: Variadic[*tuple[Unknown, ...]] +reveal_type(Variadic()) # revealed: Variadic[()] +reveal_type(Variadic(1, "a")) # revealed: Variadic[int, str] def _(i: int, s: str) -> None: reveal_type(Positional((i, s))) # revealed: Positional[int, str] - reveal_type(Variadic(i, s)) # revealed: Variadic[*tuple[Unknown, ...]] + reveal_type(Variadic(i, s)) # revealed: Variadic[int, str] ``` ### Unspecified type arguments @@ -301,8 +300,9 @@ def materialized_default[*Ts = *tuple[Any, ...]]() -> None: ### Starred variadic parameters -An unpacked `TypeVarTuple` can annotate `*args`. Inferring the `TypeVarTuple` from arguments matched -to the variadic parameter is not yet supported, so these calls use a gradual specialization. +An unpacked `TypeVarTuple` can annotate `*args`. Arguments matched to the variadic parameter infer +the complete type-variable tuple, preserving the shape of forwarded tuples and excluding ordinary +positional and keyword-only parameters. ```py def simple[*Ts](*args: *Ts) -> tuple[*Ts]: @@ -316,29 +316,61 @@ def with_kw_only[T, *Ts](*args: *Ts, kw: T) -> tuple[*Ts, T]: raise NotImplementedError def f(i: int, s: str, b: bool, t: tuple[int, str], vt: tuple[int, ...]) -> None: - reveal_type(simple()) # revealed: tuple[Unknown, ...] - reveal_type(simple(i, s)) # revealed: tuple[Unknown, ...] - reveal_type(simple(*(i, s))) # revealed: tuple[Unknown, ...] - reveal_type(simple(t)) # revealed: tuple[Unknown, ...] - reveal_type(simple(*t)) # revealed: tuple[Unknown, ...] - reveal_type(simple(*vt)) # revealed: tuple[Unknown, ...] - - reveal_type(with_prefix(i)) # revealed: tuple[int, *tuple[Unknown, ...]] - reveal_type(with_prefix(i, s, b)) # revealed: tuple[int, *tuple[Unknown, ...]] - reveal_type(with_prefix(*t)) # revealed: tuple[int, *tuple[Unknown, ...]] - reveal_type(with_prefix(i, *t)) # revealed: tuple[int, *tuple[Unknown, ...]] - reveal_type(with_prefix(*vt)) # revealed: tuple[int, *tuple[Unknown, ...]] - reveal_type(with_prefix(i, *vt)) # revealed: tuple[int, *tuple[Unknown, ...]] - - reveal_type(with_kw_only(kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] - reveal_type(with_kw_only(i, s, kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] - reveal_type(with_kw_only(t, kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] - reveal_type(with_kw_only(*t, kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] - reveal_type(with_kw_only(vt, kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] - reveal_type(with_kw_only(*vt, kw=b)) # revealed: tuple[*tuple[Unknown, ...], bool] + reveal_type(simple()) # revealed: tuple[()] + reveal_type(simple(i, s)) # revealed: tuple[int, str] + reveal_type(simple(*(i, s))) # revealed: tuple[int, str] + reveal_type(simple(t)) # revealed: tuple[tuple[int, str]] + reveal_type(simple(*t)) # revealed: tuple[int, str] + reveal_type(simple(*vt)) # revealed: tuple[int, ...] + + reveal_type(with_prefix(i)) # revealed: tuple[int] + reveal_type(with_prefix(i, s, b)) # revealed: tuple[int, str, bool] + reveal_type(with_prefix(*t)) # revealed: tuple[int, str] + reveal_type(with_prefix(i, *t)) # revealed: tuple[int, int, str] + reveal_type(with_prefix(*vt)) # revealed: tuple[int, *tuple[int, ...]] + reveal_type(with_prefix(i, *vt)) # revealed: tuple[int, *tuple[int, ...]] + + reveal_type(with_kw_only(kw=b)) # revealed: tuple[bool] + reveal_type(with_kw_only(i, s, kw=b)) # revealed: tuple[int, str, bool] + reveal_type(with_kw_only(t, kw=b)) # revealed: tuple[tuple[int, str], bool] + reveal_type(with_kw_only(*t, kw=b)) # revealed: tuple[int, str, bool] + reveal_type(with_kw_only(vt, kw=b)) # revealed: tuple[tuple[int, ...], bool] + reveal_type(with_kw_only(*vt, kw=b)) # revealed: tuple[*tuple[int, ...], bool] # error: [missing-argument] "No argument provided for required parameter `kw` of function `with_kw_only`" - reveal_type(with_kw_only(i, s, b)) # revealed: tuple[*tuple[Unknown, ...], Unknown] + reveal_type(with_kw_only(i, s, b)) # revealed: tuple[int, str, bool, Unknown] +``` + +Variadic inference must preserve contextual argument types, including contexts that contain an outer +type variable. + +```py +from typing import TypedDict + +class Payload(TypedDict): + value: int + +def contextual[T](value: T) -> None: + concrete: tuple[Payload, list[int]] = simple({"value": 1}, []) + generic: tuple[Payload, T] = simple({"value": 1}, value) +``` + +Fixed elements beside the type-variable tuple retain their individual bound diagnostics without +reporting the same error twice. + +```py +def bounded_suffix[T: str, *Ts](*args: *tuple[*Ts, T]) -> tuple[*Ts, T]: + raise NotImplementedError + +def bounded_arguments[U: bytes, T: str, *Ts](first: U, *args: *tuple[*Ts, T]) -> tuple[*Ts, T]: + raise NotImplementedError + +bounded_suffix("ok", 1) # error: [invalid-argument-type] "upper bound `str`" +bounded_arguments( + 1, # error: [invalid-argument-type] "upper bound `bytes`" + "ok", + 2, # error: [invalid-argument-type] "upper bound `str`" +) ``` ### Callable inference @@ -793,8 +825,7 @@ accept_str_in_between(True, "phase", "status", b"ok") accept_str_in_between(True, b"ok") accept_str_in_between(True, 1, b"bad") # error: [invalid-argument-type] -# TODO: Infer the `TypeVarTuple` from arguments matched to the variadic parameter. -reveal_type(remove_bytes(1, "record", b"sum")) # revealed: tuple[Unknown, ...] +reveal_type(remove_bytes(1, "record", b"sum")) # revealed: tuple[Literal[1], Literal["record"]] ``` ## `@staticmethod` and `@classmethod` diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 97119739ab5ecd..c8619de6d4013a 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -58,7 +58,7 @@ use crate::types::signatures::{ CallableSignature, Parameter, ParameterDisplayName, ParameterKind, Parameters, ParametersKind, PartialApplication, PartialSignatureApplication, }; -use crate::types::tuple::{TupleLength, TupleSpec, TupleType, VariableSegment}; +use crate::types::tuple::{TupleLength, TupleSpec, TupleSpecBuilder, TupleType, VariableSegment}; use crate::types::typed_dict::{TypedDictOpenness, extract_unpacked_typed_dict_from_value_type}; use crate::types::typevar::{BoundTypeVarIdentity, TypeVarNonceGenerator, TypeVarSet}; use crate::types::visitor::{ @@ -5920,6 +5920,175 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> { self.inference = Some(inference); } + /// Infer a variadic type-variable tuple from all positional arguments bound to `*args`. + /// + /// A type-variable tuple represents the complete argument sequence, not an independent type + /// variable for each argument. The ordinary tuple inference machinery also accounts for any + /// fixed elements surrounding the pack: + /// + /// ```python + /// def collect[*Ts](*args: *Ts) -> tuple[*Ts]: ... + /// def discard_suffix[*Ts](*args: *tuple[*Ts, bytes]) -> tuple[*Ts]: ... + /// + /// collect(1, "two") + /// discard_suffix(1, "two", b"last") + /// ``` + fn infer_typevartuple_argument_constraints<'c>( + &self, + builder: &mut SpecializationBuilder<'db, 'c>, + ) -> Result<(), SpecializationError<'db>> { + let db = self.db; + let Some((parameter_index, parameter)) = self.signature.parameters().variadic() else { + return Ok(()); + }; + if !parameter.has_starred_annotation() { + return Ok(()); + } + + let (formal, typevartuple) = match parameter.annotated_type() { + Type::TypeVar(typevar) if typevar.is_typevartuple(db) => ( + Type::tuple(Some(TupleType::unpacked_typevartuple( + db, self.env, typevar, + ))), + typevar, + ), + annotation => { + let Some(typevartuple) = + annotation.exact_tuple_instance_spec(db).and_then(|tuple| { + match tuple.as_ref() { + TupleSpec::Variable(variable) => variable.variable().typevartuple(), + TupleSpec::Fixed(_) => None, + } + }) + else { + return Ok(()); + }; + (annotation, typevartuple) + } + }; + + let contains_typevartuple = |ty| { + any_over_type(db, self.env, ty, true, |nested| { + matches!( + nested, + Type::TypeVar(typevar) + if typevar.identity(db) == typevartuple.identity(db) + ) + }) + }; + + // Other occurrences of the same pack require joint inference. Preserve the existing + // behavior until those constraints can be solved together instead of independently + // widening an already established specialization. + if !contains_typevartuple(self.return_ty) + || self + .enumerate_argument_types() + .any(|(argument_index, _, argument, _)| { + !matches!(argument, Argument::Synthetic) + && self.argument_matches[argument_index].iter().any(|matched| { + matched.index != parameter_index + && contains_typevartuple( + self.signature.parameters()[matched.index].annotated_type(), + ) + }) + }) + { + return Ok(()); + } + + let mut actual = TupleSpecBuilder::with_capacity(self.arguments.len()); + for (argument_index, _, argument, argument_types) in self.enumerate_argument_types() { + let matches = &self.argument_matches[argument_index]; + if !matches + .iter() + .any(|matched| matched.index == parameter_index) + { + continue; + } + + if matches!(argument, Argument::Variadic) { + let Some(argument_type) = argument_types.get_default() else { + return Ok(()); + }; + // Iteration would merge union branches, losing correlations between their + // argument counts and element types. + if matches!(argument_type.resolve_type_alias(db), Type::Union(_)) { + return Ok(()); + } + let mut argument_tuple = argument_type.iterate(db, self.env).into_owned(); + let consumed_prefix = matches + .parameters + .iter() + .take_while(|matched| matched.index != parameter_index) + .count(); + if consumed_prefix != 0 { + let Ok(consumed_prefix) = i32::try_from(consumed_prefix) else { + return Ok(()); + }; + let Ok(sliced) = argument_tuple.py_slice_type( + db, + self.env, + Some(consumed_prefix), + None, + None, + ) else { + return Ok(()); + }; + let Some(sliced) = sliced.exact_tuple_instance_spec(db) else { + return Ok(()); + }; + argument_tuple = sliced.into_owned(); + } + actual = actual.concat(db, self.env, &argument_tuple); + continue; + } + + for matched in matches + .iter() + .filter(|matched| matched.index == parameter_index) + { + let declared_type = matched + .expected_type + .unwrap_or_else(|| parameter.annotated_type()); + actual.push( + matched + .argument_type + .unwrap_or_else(|| argument_types.get_for_declared_type(declared_type)), + ); + } + } + + let actual = Type::tuple(TupleType::new(db, self.env, &actual.build())); + + // An inference fixpoint can supply its previous result as context. Only defer to that + // context when independently inferring the argument pack would actually contradict it. + if let Some(expected) = self.call_expression_tcx.annotation + && !expected.is_dynamic() + && let Some(generic_context) = self.signature.generic_context + { + let constraints = ConstraintSetBuilder::new(); + let mut projected = + SpecializationBuilder::new(db, self.env, &constraints, self.inferable_typevars); + projected.infer(formal, actual)?; + + if let Ok(inference) = projected.build_inference_with(generic_context, |_, _| None) { + let candidate = self + .return_ty + .apply_specialization(db, inference.specialization(db)); + + if !candidate.is_assignable_to(db, self.env, expected) + && !candidate + .promote(db, self.env) + .is_assignable_to(db, self.env, expected) + { + return Ok(()); + } + } + } + + builder.infer(formal, actual) + } + fn infer_argument_constraints<'c>( &mut self, builder: &mut SpecializationBuilder<'db, 'c>, @@ -5936,8 +6105,8 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> { let parameter_index = matched_parameter.index; let parameter = ¶meters[parameter_index]; let parameter_type = parameter.annotated_type(); - // TODO: Infer a `TypeVarTuple` from all matched positional arguments as a single - // tuple. Fixed elements beside that pack can still infer ordinary type variables. + // A `TypeVarTuple` is inferred once from its complete argument tuple below. Fixed + // elements beside that pack can still infer ordinary type variables individually. if parameter.has_starred_annotation() && matched_parameter.expected_type.is_none() && (matches!( @@ -5975,6 +6144,23 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> { } } + if let Err(error) = self.infer_typevartuple_argument_constraints(builder) + && !specialization_errors.iter().any(|existing| { + matches!( + existing, + BindingError::SpecializationError { + error: existing, + .. + } if existing == &error + ) + }) + { + specialization_errors.push(BindingError::SpecializationError { + error, + argument_index: None, + }); + } + preferred_type_mappings .iter() .all(|(&identity, &preferred_ty)| {