diff --git a/crates/ty_python_semantic/resources/mdtest/loops/for.md b/crates/ty_python_semantic/resources/mdtest/loops/for.md index 1d657575c6..3916fa884a 100644 --- a/crates/ty_python_semantic/resources/mdtest/loops/for.md +++ b/crates/ty_python_semantic/resources/mdtest/loops/for.md @@ -401,6 +401,66 @@ def _(x: Intersection[NotIterable1, NotIterable2]): reveal_type(item) # revealed: Unknown ``` +## Intersection of fixed-length tuples + +When iterating over an intersection of two fixed-length tuples with the same length, we should +intersect the element types position-by-position. + +```py +from ty_extensions import Intersection + +def _(x: Intersection[tuple[int, str], tuple[object, object]]): + # `tuple[int, str]` yields `int | str` when iterated. + # `tuple[object, object]` yields `object` when iterated. + # The intersection should yield `(int & object) | (str & object)` = `int | str`. + for item in x: + reveal_type(item) # revealed: int | str +``` + +## Intersection of fixed-length tuple with homogeneous iterable + +When iterating over an intersection of a fixed-length tuple with a class that implements `__iter__` +returning a homogeneous iterator, we should preserve the fixed-length structure and intersect each +element type with the iterator's element type. + +```py +from collections.abc import Iterator + +class Foo: + def __iter__(self) -> Iterator[object]: + raise NotImplementedError + +def _(x: tuple[int, str, bytes]): + if isinstance(x, Foo): + # The intersection `tuple[int, str, bytes] & Foo` should iterate as + # `tuple[int & object, str & object, bytes & object]` = `tuple[int, str, bytes]` + a, b, c = x + reveal_type(a) # revealed: int + reveal_type(b) # revealed: str + reveal_type(c) # revealed: bytes + reveal_type(tuple(x)) # revealed: tuple[int, str, bytes] +``` + +## Intersection of homogeneous iterables + +When iterating over an intersection of two types that both yield homogeneous variable-length tuple +specs, we should intersect their element types. + +```py +from collections.abc import Iterator + +class Foo: + def __iter__(self) -> Iterator[object]: + raise NotImplementedError + +def _(x: list[int]): + if isinstance(x, Foo): + # `list[int]` yields `int`, `Foo` yields `object`. + # The intersection should yield `int & object` = `int`. + for item in x: + reveal_type(item) # revealed: int +``` + ## Possibly-not-callable `__iter__` method ```py diff --git a/crates/ty_python_semantic/src/types/tuple.rs b/crates/ty_python_semantic/src/types/tuple.rs index e551871539..4d6611e4cb 100644 --- a/crates/ty_python_semantic/src/types/tuple.rs +++ b/crates/ty_python_semantic/src/types/tuple.rs @@ -1793,13 +1793,16 @@ impl<'db> TupleSpecBuilder<'db> { /// tuple-spec for `tuple[object, object]`, the result will be a tuple-spec builder for /// `tuple[int, str]` (since `int & object` simplifies to `int`, and `str & object` to `str`). /// - /// To keep things simple, we currently only attempt to preserve the "fixed-length-ness" of - /// a tuple spec if both `self` and `other` have the exact same length. For example, - /// if `self` is a tuple-spec builder for `tuple[int, str]` and `other` is a tuple-spec for - /// `tuple[int, str, bytes]`, the result will be a tuple-spec builder for - /// `tuple[int & str & bytes, ...]`. + /// We preserve "fixed-length-ness" in the following cases: + /// - Both `self` and `other` are fixed-length with the same length: element-wise intersection + /// - `self` is fixed-length and `other` is a homogeneous variable-length tuple: intersect each + /// element of `self` with `other`'s variable type + /// + /// For other cases (e.g., different fixed lengths, or complex variable-length tuples with + /// prefixes or suffixes), we fall back to a homogeneous variable-length tuple. pub(crate) fn intersect(mut self, db: &'db dyn Db, other: &TupleSpec<'db>) -> Self { match (&mut self, other) { + // Both fixed-length with the same length: element-wise intersection. (TupleSpecBuilder::Fixed(our_elements), TupleSpec::Fixed(new_elements)) if our_elements.len() == new_elements.len() => { @@ -1809,6 +1812,39 @@ impl<'db> TupleSpecBuilder<'db> { self } + // Fixed-length intersected with a homogeneous variable-length tuple: + // intersect each element with the variable type. + (TupleSpecBuilder::Fixed(our_elements), TupleSpec::Variable(var)) + if var.prefix.is_empty() && var.suffix.is_empty() => + { + for existing in our_elements.iter_mut() { + *existing = IntersectionType::from_elements(db, [*existing, var.variable]); + } + self + } + + // Variable-length intersected with a homogeneous variable-length tuple: + // intersect prefix, variable, and suffix with the variable type + ( + TupleSpecBuilder::Variable { + prefix, + variable, + suffix, + }, + TupleSpec::Variable(other_var), + ) if other_var.prefix.is_empty() && other_var.suffix.is_empty() => { + for existing in prefix.iter_mut() { + *existing = + IntersectionType::from_elements(db, [*existing, other_var.variable]); + } + *variable = IntersectionType::from_elements(db, [*variable, other_var.variable]); + for existing in suffix.iter_mut() { + *existing = + IntersectionType::from_elements(db, [*existing, other_var.variable]); + } + self + } + _ => { let intersected = IntersectionType::from_elements( db,