[red-knot] Improve type inference for iteration over heterogenous tuples (#13314)

Followup to #13295
This commit is contained in:
Alex Waygood 2024-09-10 15:13:50 -04:00 committed by GitHub
parent a528edad35
commit 2ca78721e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 34 additions and 0 deletions

View file

@ -408,6 +408,18 @@ impl<'db> Type<'db> {
/// pass
/// ```
fn iterate(&self, db: &'db dyn Db) -> IterationOutcome<'db> {
if let Type::Tuple(tuple_type) = self {
return IterationOutcome::Iterable {
element_ty: tuple_type
.elements(db)
.iter()
.fold(UnionBuilder::new(db), |builder, element| {
builder.add(*element)
})
.build(),
};
}
// `self` represents the type of the iterable;
// `__iter__` and `__next__` are both looked up on the class of the iterable:
let iterable_meta_type = self.to_meta_type(db);

View file

@ -4355,6 +4355,28 @@ mod tests {
Ok(())
}
#[test]
fn for_loop_with_heterogenous_tuple() -> anyhow::Result<()> {
let mut db = setup_db();
db.write_dedented(
"src/a.py",
"
for x in (1, 'a', b'foo'):
pass
",
)?;
assert_public_ty(
&db,
"src/a.py",
"x",
r#"Literal[1] | Literal["a"] | Literal[b"foo"]"#,
);
Ok(())
}
#[test]
fn except_handler_single_exception() -> anyhow::Result<()> {
let mut db = setup_db();