Respect intersections in iterations

This commit is contained in:
Charlie Marsh 2025-12-13 14:26:55 -05:00
parent ff0ed4e752
commit b6a99c9acf
3 changed files with 74 additions and 3 deletions

View File

@ -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

View File

@ -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(_)

View File

@ -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) => {