Skip to content

Commit

Permalink
[red-knot] Add narrowing for 'while' loops (#14947)
Browse files Browse the repository at this point in the history
## Summary

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

closes #14861 

## Test Plan

New Markdown tests.
  • Loading branch information
sharkdp authored Dec 13, 2024
1 parent be4ce16 commit d7ce548
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 1 deletion.
58 changes: 58 additions & 0 deletions crates/red_knot_python_semantic/resources/mdtest/narrow/while.md
Original file line number Diff line number Diff line change
@@ -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()
```
2 changes: 2 additions & 0 deletions crates/red_knot_python_semantic/src/semantic_index/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down

0 comments on commit d7ce548

Please sign in to comment.