diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md new file mode 100644 index 0000000000..ac91a576c9 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/while.md @@ -0,0 +1,58 @@ +# Narrowing in `while` loops + +We only make sure that narrowing works for `while` loops in general, we do not exhaustively test all +narrowing forms here, as they are covered in other tests. + +Note how type narrowing works subtly different from `if` ... `else`, because the negated constraint +is retained after the loop. + +## Basic `while` loop + +```py +def next_item() -> int | None: ... + +x = next_item() + +while x is not None: + reveal_type(x) # revealed: int + x = next_item() + +reveal_type(x) # revealed: None +``` + +## `while` loop with `else` + +```py +def next_item() -> int | None: ... + +x = next_item() + +while x is not None: + reveal_type(x) # revealed: int + x = next_item() +else: + reveal_type(x) # revealed: None + +reveal_type(x) # revealed: None +``` + +## Nested `while` loops + +```py +from typing import Literal + +def next_item() -> Literal[1, 2, 3]: ... + +x = next_item() + +while x != 1: + reveal_type(x) # revealed: Literal[2, 3] + + while x != 2: + # TODO: this should be Literal[1, 3]; Literal[3] is only correct + # in the first loop iteration + reveal_type(x) # revealed: Literal[3] + x = next_item() + + x = next_item() +``` 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 c75af214d8..27e657aba9 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs @@ -833,6 +833,7 @@ where self.visit_expr(test); let pre_loop = self.flow_snapshot(); + let constraint = self.record_expression_constraint(test); // Save aside any break states from an outer loop let saved_break_states = std::mem::take(&mut self.loop_break_states); @@ -852,6 +853,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.record_negated_constraint(constraint); self.visit_body(orelse); // Breaking out of a while loop bypasses the `else` clause, so merge in the break diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 6b4393f588..1b732ce531 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -2092,7 +2092,7 @@ impl<'db> TypeInferenceBuilder<'db> { orelse, } = while_statement; - self.infer_expression(test); + self.infer_standalone_expression(test); self.infer_body(body); self.infer_body(orelse); }