[red-knot] Add narrowing for 'while' loops (#14947)
Some checks are pending
CI / cargo fmt (push) Waiting to run
CI / cargo build (release) (push) Waiting to run
CI / cargo shear (push) Blocked by required conditions
CI / Determine changes (push) Waiting to run
CI / cargo clippy (push) Blocked by required conditions
CI / cargo test (linux) (push) Blocked by required conditions
CI / cargo test (linux, release) (push) Blocked by required conditions
CI / cargo test (windows) (push) Blocked by required conditions
CI / cargo test (wasm) (push) Blocked by required conditions
CI / cargo build (msrv) (push) Blocked by required conditions
CI / cargo fuzz build (push) Blocked by required conditions
CI / fuzz parser (push) Blocked by required conditions
CI / test scripts (push) Blocked by required conditions
CI / ecosystem (push) Blocked by required conditions
CI / python package (push) Waiting to run
CI / pre-commit (push) Waiting to run
CI / mkdocs (push) Waiting to run
CI / formatter instabilities and black similarity (push) Blocked by required conditions
CI / test ruff-lsp (push) Blocked by required conditions
CI / benchmarks (push) Blocked by required conditions

## Summary

Add type narrowing for `while` loops and corresponding `else` branches.

closes #14861 

## Test Plan

New Markdown tests.
This commit is contained in:
David Peter 2024-12-13 07:40:14 +01:00 committed by GitHub
parent be4ce16735
commit d7ce548893
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 61 additions and 1 deletions

View file

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

View file

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

View file

@ -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);
}