From b93d0ab57c86f2f48b52b56d98f0fc2e5f5e7d10 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Tue, 10 Sep 2024 18:04:35 -0400 Subject: [PATCH] [red-knot] Add control flow for `for` loops (#13318) --- .../src/semantic_index/builder.rs | 39 +++++++- .../src/types/infer.rs | 95 ++++++++++++++++++- 2 files changed, 124 insertions(+), 10 deletions(-) diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs index 3f440a89b3..8fb5a7f941 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -575,14 +575,23 @@ where self.flow_merge(pre_if); } } - ast::Stmt::While(node) => { - self.visit_expr(&node.test); + ast::Stmt::While(ast::StmtWhile { + test, + body, + orelse, + range: _, + }) => { + self.visit_expr(test); let pre_loop = self.flow_snapshot(); // Save aside any break states from an outer loop let saved_break_states = std::mem::take(&mut self.loop_break_states); - self.visit_body(&node.body); + + // TODO: definitions created inside the body should be fully visible + // to other statements/expressions inside the body --Alex/Carl + self.visit_body(body); + // Get the break states from the body of this loop, and restore the saved outer // ones. let break_states = @@ -591,7 +600,7 @@ where // We may execute the `else` clause without ever executing the body, so merge in // the pre-loop state before visiting `else`. self.flow_merge(pre_loop); - self.visit_body(&node.orelse); + self.visit_body(orelse); // Breaking out of a while loop bypasses the `else` clause, so merge in the break // states after visiting `else`. @@ -625,15 +634,35 @@ where orelse, }, ) => { - // TODO add control flow similar to `ast::Stmt::While` above self.add_standalone_expression(iter); self.visit_expr(iter); + + let pre_loop = self.flow_snapshot(); + let saved_break_states = std::mem::take(&mut self.loop_break_states); + debug_assert!(self.current_assignment.is_none()); self.current_assignment = Some(for_stmt.into()); self.visit_expr(target); self.current_assignment = None; + + // TODO: Definitions created by loop variables + // (and definitions created inside the body) + // are fully visible to other statements/expressions inside the body --Alex/Carl self.visit_body(body); + + let break_states = + std::mem::replace(&mut self.loop_break_states, saved_break_states); + + // We may execute the `else` clause without ever executing the body, so merge in + // the pre-loop state before visiting `else`. + self.flow_merge(pre_loop); self.visit_body(orelse); + + // Breaking out of a `for` loop bypasses the `else` clause, so merge in the break + // states after visiting `else`. + for break_state in break_states { + self.flow_merge(break_state); + } } ast::Stmt::Match(ast::StmtMatch { subject, diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index d9b88dddb5..3dd1a4d1ff 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -4271,7 +4271,92 @@ mod tests { ", )?; - assert_public_ty(&db, "src/a.py", "x", "int"); + assert_public_ty(&db, "src/a.py", "x", "Unbound | int"); + + Ok(()) + } + + #[test] + fn for_loop_with_previous_definition() -> anyhow::Result<()> { + let mut db = setup_db(); + + db.write_dedented( + "src/a.py", + " + class IntIterator: + def __next__(self) -> int: + return 42 + + class IntIterable: + def __iter__(self) -> IntIterator: + return IntIterator() + + x = 'foo' + + for x in IntIterable(): + pass + ", + )?; + + assert_public_ty(&db, "src/a.py", "x", r#"Literal["foo"] | int"#); + + Ok(()) + } + + #[test] + fn for_loop_no_break() -> anyhow::Result<()> { + let mut db = setup_db(); + + db.write_dedented( + "src/a.py", + " + class IntIterator: + def __next__(self) -> int: + return 42 + + class IntIterable: + def __iter__(self) -> IntIterator: + return IntIterator() + + for x in IntIterable(): + pass + else: + x = 'foo' + ", + )?; + + // The `for` loop can never break, so the `else` clause will always be executed, + // meaning that the visible definition by the end of the scope is solely determined + // by the `else` clause + assert_public_ty(&db, "src/a.py", "x", r#"Literal["foo"]"#); + + Ok(()) + } + + #[test] + fn for_loop_may_break() -> anyhow::Result<()> { + let mut db = setup_db(); + + db.write_dedented( + "src/a.py", + " + class IntIterator: + def __next__(self) -> int: + return 42 + + class IntIterable: + def __iter__(self) -> IntIterator: + return IntIterator() + + for x in IntIterable(): + if x > 5: + break + else: + x = 'foo' + ", + )?; + + assert_public_ty(&db, "src/a.py", "x", r#"int | Literal["foo"]"#); Ok(()) } @@ -4292,7 +4377,7 @@ mod tests { ", )?; - assert_public_ty(&db, "src/a.py", "x", "int"); + assert_public_ty(&db, "src/a.py", "x", "Unbound | int"); Ok(()) } @@ -4320,7 +4405,7 @@ mod tests { ", )?; - assert_scope_ty(&db, "src/a.py", &["foo"], "x", "Unknown"); + assert_scope_ty(&db, "src/a.py", &["foo"], "x", "Unbound | Unknown"); Ok(()) } @@ -4347,7 +4432,7 @@ mod tests { )?; // TODO(Alex) async iterables/iterators! - assert_scope_ty(&db, "src/a.py", &["foo"], "x", "Unknown"); + assert_scope_ty(&db, "src/a.py", &["foo"], "x", "Unbound | Unknown"); Ok(()) } @@ -4368,7 +4453,7 @@ mod tests { &db, "src/a.py", "x", - r#"Literal[1] | Literal["a"] | Literal[b"foo"]"#, + r#"Unbound | Literal[1] | Literal["a"] | Literal[b"foo"]"#, ); Ok(())