From ec3d5ebda2a22d6ddf52af3799bbe365e9e177c8 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Wed, 30 Jul 2025 17:12:21 +0100 Subject: [PATCH] [ty] Upcast heterogeneous and mixed tuples to homogeneous tuples where it's necessary to solve a `TypeVar` (#19635) ## Summary This PR improves our generics solver such that we are able to solve the `TypeVar` in this snippet to `int | str` (the union of the elements in the heterogeneous tuple) by upcasting the heterogeneous tuple to its pure-homogeneous-tuple supertype: ```py def f[T](x: tuple[T, ...]) -> T: return x[0] def g(x: tuple[int, str]): reveal_type(f(x)) ``` ## Test Plan Mdtests. Some TODOs remain in the mdtest regarding solving `TypeVar`s for mixed tuples, but I think this PR on its own is a significant step forward for our generics solver when it comes to tuple types. --------- Co-authored-by: Douglas Creager --- .../mdtest/generics/legacy/functions.md | 25 ++++++++++------- .../mdtest/generics/pep695/functions.md | 25 ++++++++++------- .../ty_python_semantic/src/types/generics.rs | 24 +++++++++-------- crates/ty_python_semantic/src/types/tuple.rs | 27 +++++++++++++++++++ 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md index da947217f8..65872213f2 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md @@ -145,27 +145,34 @@ T = TypeVar("T") def takes_mixed_tuple_suffix(x: tuple[int, bytes, *tuple[str, ...], T, int]) -> T: return x[-2] -# TODO: revealed: Literal[True] -reveal_type(takes_mixed_tuple_suffix((1, b"foo", "bar", "baz", True, 42))) # revealed: Unknown - def takes_mixed_tuple_prefix(x: tuple[int, T, *tuple[str, ...], bool, int]) -> T: return x[1] -# TODO: revealed: Literal[b"foo"] -reveal_type(takes_mixed_tuple_prefix((1, b"foo", "bar", "baz", True, 42))) # revealed: Unknown +def _(x: tuple[int, bytes, *tuple[str, ...], bool, int]): + reveal_type(takes_mixed_tuple_suffix(x)) # revealed: bool + reveal_type(takes_mixed_tuple_prefix(x)) # revealed: bytes + +reveal_type(takes_mixed_tuple_suffix((1, b"foo", "bar", "baz", True, 42))) # revealed: Literal[True] +reveal_type(takes_mixed_tuple_prefix((1, b"foo", "bar", "baz", True, 42))) # revealed: Literal[b"foo"] def takes_fixed_tuple(x: tuple[T, int]) -> T: return x[0] +def _(x: tuple[str, int]): + reveal_type(takes_fixed_tuple(x)) # revealed: str + reveal_type(takes_fixed_tuple((True, 42))) # revealed: Literal[True] def takes_homogeneous_tuple(x: tuple[T, ...]) -> T: return x[0] -# TODO: revealed: Literal[42] -reveal_type(takes_homogeneous_tuple((42,))) # revealed: Unknown -# TODO: revealed: Literal[42, 43] -reveal_type(takes_homogeneous_tuple((42, 43))) # revealed: Unknown +def _(x: tuple[str, int], y: tuple[bool, ...], z: tuple[int, str, *tuple[range, ...], bytes]): + reveal_type(takes_homogeneous_tuple(x)) # revealed: str | int + reveal_type(takes_homogeneous_tuple(y)) # revealed: bool + reveal_type(takes_homogeneous_tuple(z)) # revealed: int | str | range | bytes + +reveal_type(takes_homogeneous_tuple((42,))) # revealed: Literal[42] +reveal_type(takes_homogeneous_tuple((42, 43))) # revealed: Literal[42, 43] ``` ## Inferring a bound typevar diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md index bc7f96c0e4..aad1bb55b9 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md @@ -131,27 +131,34 @@ reveal_type(takes_in_protocol(ExplicitGenericSub[str]())) # revealed: str def takes_mixed_tuple_suffix[T](x: tuple[int, bytes, *tuple[str, ...], T, int]) -> T: return x[-2] -# TODO: revealed: Literal[True] -reveal_type(takes_mixed_tuple_suffix((1, b"foo", "bar", "baz", True, 42))) # revealed: Unknown - def takes_mixed_tuple_prefix[T](x: tuple[int, T, *tuple[str, ...], bool, int]) -> T: return x[1] -# TODO: revealed: Literal[b"foo"] -reveal_type(takes_mixed_tuple_prefix((1, b"foo", "bar", "baz", True, 42))) # revealed: Unknown +def _(x: tuple[int, bytes, *tuple[str, ...], bool, int]): + reveal_type(takes_mixed_tuple_suffix(x)) # revealed: bool + reveal_type(takes_mixed_tuple_prefix(x)) # revealed: bytes + +reveal_type(takes_mixed_tuple_suffix((1, b"foo", "bar", "baz", True, 42))) # revealed: Literal[True] +reveal_type(takes_mixed_tuple_prefix((1, b"foo", "bar", "baz", True, 42))) # revealed: Literal[b"foo"] def takes_fixed_tuple[T](x: tuple[T, int]) -> T: return x[0] +def _(x: tuple[str, int]): + reveal_type(takes_fixed_tuple(x)) # revealed: str + reveal_type(takes_fixed_tuple((True, 42))) # revealed: Literal[True] def takes_homogeneous_tuple[T](x: tuple[T, ...]) -> T: return x[0] -# TODO: revealed: Literal[42] -reveal_type(takes_homogeneous_tuple((42,))) # revealed: Unknown -# TODO: revealed: Literal[42, 43] -reveal_type(takes_homogeneous_tuple((42, 43))) # revealed: Unknown +def _(x: tuple[str, int], y: tuple[bool, ...], z: tuple[int, str, *tuple[range, ...], bytes]): + reveal_type(takes_homogeneous_tuple(x)) # revealed: str | int + reveal_type(takes_homogeneous_tuple(y)) # revealed: bool + reveal_type(takes_homogeneous_tuple(z)) # revealed: int | str | range | bytes + +reveal_type(takes_homogeneous_tuple((42,))) # revealed: Literal[42] +reveal_type(takes_homogeneous_tuple((42, 43))) # revealed: Literal[42, 43] ``` ## Inferring a bound typevar diff --git a/crates/ty_python_semantic/src/types/generics.rs b/crates/ty_python_semantic/src/types/generics.rs index f19a1a477e..f765bfdc66 100644 --- a/crates/ty_python_semantic/src/types/generics.rs +++ b/crates/ty_python_semantic/src/types/generics.rs @@ -761,17 +761,19 @@ impl<'db> SpecializationBuilder<'db> { (Type::Tuple(formal_tuple), Type::Tuple(actual_tuple)) => { let formal_tuple = formal_tuple.tuple(self.db); let actual_tuple = actual_tuple.tuple(self.db); - match (formal_tuple, actual_tuple) { - (TupleSpec::Fixed(formal_tuple), TupleSpec::Fixed(actual_tuple)) => { - if formal_tuple.len() == actual_tuple.len() { - for (formal_element, actual_element) in formal_tuple.elements().zip(actual_tuple.elements()) { - self.infer(*formal_element, *actual_element)?; - } - } - } - - // TODO: Infer specializations of variable-length tuples - (TupleSpec::Variable(_), _) | (_, TupleSpec::Variable(_)) => {} + let Some(most_precise_length) = formal_tuple.len().most_precise(actual_tuple.len()) else { + return Ok(()); + }; + let Ok(formal_tuple) = formal_tuple.resize(self.db, most_precise_length) else { + return Ok(()); + }; + let Ok(actual_tuple) = actual_tuple.resize(self.db, most_precise_length) else { + return Ok(()); + }; + for (formal_element, actual_element) in + formal_tuple.all_elements().zip(actual_tuple.all_elements()) + { + self.infer(*formal_element, *actual_element)?; } } diff --git a/crates/ty_python_semantic/src/types/tuple.rs b/crates/ty_python_semantic/src/types/tuple.rs index 268bd42d1e..cd2f1eb026 100644 --- a/crates/ty_python_semantic/src/types/tuple.rs +++ b/crates/ty_python_semantic/src/types/tuple.rs @@ -71,6 +71,33 @@ impl TupleLength { } } + /// Given two [`TupleLength`]s, return the more precise instance, + /// if it makes sense to consider one more precise than the other. + pub(crate) fn most_precise(self, other: Self) -> Option { + match (self, other) { + // A fixed-length tuple is equally as precise as another fixed-length tuple if they + // have the same length. For two differently sized fixed-length tuples, however, + // neither tuple length is more precise than the other: the two tuple lengths are + // entirely disjoint. + (TupleLength::Fixed(left), TupleLength::Fixed(right)) => { + (left == right).then_some(self) + } + + // A fixed-length tuple is more precise than a variable-length one. + (fixed @ TupleLength::Fixed(_), TupleLength::Variable(..)) + | (TupleLength::Variable(..), fixed @ TupleLength::Fixed(_)) => Some(fixed), + + // For two variable-length tuples, the tuple with the larger number + // of required items is more precise. + (TupleLength::Variable(..), TupleLength::Variable(..)) => { + Some(match self.minimum().cmp(&other.minimum()) { + Ordering::Less => other, + Ordering::Equal | Ordering::Greater => self, + }) + } + } + } + pub(crate) fn display_minimum(self) -> String { let minimum_length = self.minimum(); match self {