mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 05:45:24 +00:00
[red-knot] Add narrowing for 'while' loops (#14947)
## Summary Add type narrowing for `while` loops and corresponding `else` branches. closes #14861 ## Test Plan New Markdown tests.
This commit is contained in:
parent
be4ce16735
commit
d7ce548893
3 changed files with 61 additions and 1 deletions
|
@ -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()
|
||||||
|
```
|
|
@ -833,6 +833,7 @@ where
|
||||||
self.visit_expr(test);
|
self.visit_expr(test);
|
||||||
|
|
||||||
let pre_loop = self.flow_snapshot();
|
let pre_loop = self.flow_snapshot();
|
||||||
|
let constraint = self.record_expression_constraint(test);
|
||||||
|
|
||||||
// Save aside any break states from an outer loop
|
// Save aside any break states from an outer loop
|
||||||
let saved_break_states = std::mem::take(&mut self.loop_break_states);
|
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
|
// We may execute the `else` clause without ever executing the body, so merge in
|
||||||
// the pre-loop state before visiting `else`.
|
// the pre-loop state before visiting `else`.
|
||||||
self.flow_merge(pre_loop);
|
self.flow_merge(pre_loop);
|
||||||
|
self.record_negated_constraint(constraint);
|
||||||
self.visit_body(orelse);
|
self.visit_body(orelse);
|
||||||
|
|
||||||
// Breaking out of a while loop bypasses the `else` clause, so merge in the break
|
// Breaking out of a while loop bypasses the `else` clause, so merge in the break
|
||||||
|
|
|
@ -2092,7 +2092,7 @@ impl<'db> TypeInferenceBuilder<'db> {
|
||||||
orelse,
|
orelse,
|
||||||
} = while_statement;
|
} = while_statement;
|
||||||
|
|
||||||
self.infer_expression(test);
|
self.infer_standalone_expression(test);
|
||||||
self.infer_body(body);
|
self.infer_body(body);
|
||||||
self.infer_body(orelse);
|
self.infer_body(orelse);
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue