diff --git a/crates/ty_python_semantic/resources/mdtest/loops/for.md b/crates/ty_python_semantic/resources/mdtest/loops/for.md index ed51e51c56..0da97ff5c0 100644 --- a/crates/ty_python_semantic/resources/mdtest/loops/for.md +++ b/crates/ty_python_semantic/resources/mdtest/loops/for.md @@ -337,6 +337,31 @@ for x in Test(): reveal_type(x) # revealed: int ``` +## Intersection type via isinstance narrowing + +When we have an intersection type via `isinstance` narrowing, we should be able to infer the +iterable element type precisely: + +```py +from typing import Sequence + +def _(x: Sequence[int], y: object): + reveal_type(x) # revealed: Sequence[int] + for item in x: + reveal_type(item) # revealed: int + + if isinstance(y, list): + reveal_type(y) # revealed: Top[list[Unknown]] + for item in y: + reveal_type(item) # revealed: object + + if isinstance(x, list): + reveal_type(x) # revealed: Sequence[int] & Top[list[Unknown]] + for item in x: + # int & object simplifies to int + reveal_type(item) # revealed: int +``` + ## Possibly-not-callable `__iter__` method ```py diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 1bbcc038cc..6bcc57f03f 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -6612,6 +6612,17 @@ impl<'db> Type<'db> { None } } + Type::Intersection(intersection) => { + // For intersections, we iterate over each positive element and intersect + // the resulting element types. Negative elements don't affect iteration. + let mut elements_iter = intersection.positive_elements_or_object(db); + let first_element_spec = elements_iter.next()?.try_iterate_with_mode(db, EvaluationMode::Sync).ok()?; + let mut builder = TupleSpecBuilder::from(&*first_element_spec); + for element in elements_iter { + builder = builder.intersect(db, &*element.try_iterate_with_mode(db, EvaluationMode::Sync).ok()?); + } + Some(Cow::Owned(builder.build())) + } // N.B. These special cases aren't strictly necessary, they're just obvious optimizations Type::LiteralString | Type::Dynamic(_) => Some(Cow::Owned(TupleSpec::homogeneous(ty))), @@ -6634,7 +6645,6 @@ impl<'db> Type<'db> { | Type::SpecialForm(_) | Type::KnownInstance(_) | Type::PropertyInstance(_) - | Type::Intersection(_) | Type::AlwaysTruthy | Type::AlwaysFalsy | Type::IntLiteral(_) diff --git a/crates/ty_python_semantic/src/types/tuple.rs b/crates/ty_python_semantic/src/types/tuple.rs index f0cc3bedc8..e551871539 100644 --- a/crates/ty_python_semantic/src/types/tuple.rs +++ b/crates/ty_python_semantic/src/types/tuple.rs @@ -29,8 +29,8 @@ use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension}; use crate::types::generics::InferableTypeVars; use crate::types::{ ApplyTypeMappingVisitor, BoundTypeVarInstance, FindLegacyTypeVarsVisitor, HasRelationToVisitor, - IsDisjointVisitor, IsEquivalentVisitor, NormalizedVisitor, Type, TypeMapping, TypeRelation, - UnionBuilder, UnionType, + IntersectionType, IsDisjointVisitor, IsEquivalentVisitor, NormalizedVisitor, Type, TypeMapping, + TypeRelation, UnionBuilder, UnionType, }; use crate::types::{Truthiness, TypeContext}; use crate::{Db, FxOrderSet, Program}; @@ -1787,6 +1787,42 @@ impl<'db> TupleSpecBuilder<'db> { } } + /// Return a new tuple-spec builder that reflects the intersection of this tuple and another tuple. + /// + /// For example, if `self` is a tuple-spec builder for `tuple[int, str]` and `other` is a + /// 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, ...]`. + pub(crate) fn intersect(mut self, db: &'db dyn Db, other: &TupleSpec<'db>) -> Self { + match (&mut self, other) { + (TupleSpecBuilder::Fixed(our_elements), TupleSpec::Fixed(new_elements)) + if our_elements.len() == new_elements.len() => + { + for (existing, new) in our_elements.iter_mut().zip(new_elements.elements()) { + *existing = IntersectionType::from_elements(db, [*existing, *new]); + } + self + } + + _ => { + let intersected = IntersectionType::from_elements( + db, + self.all_elements().chain(other.all_elements()), + ); + TupleSpecBuilder::Variable { + prefix: vec![], + variable: intersected, + suffix: vec![], + } + } + } + } + pub(super) fn build(self) -> TupleSpec<'db> { match self { TupleSpecBuilder::Fixed(elements) => {