Skip to content

Commit 3f6c65e

Browse files
sransaracarljmAlexWaygood
authored
[red-knot] Fix merged type after if-else without explicit else branch (#14621)
## Summary Closes: #14593 The final type of a variable after if-statement without explicit else branch should be similar to having an explicit else branch. ## Test Plan Originally failed test cases from the bug are added. --------- Co-authored-by: Carl Meyer <[email protected]> Co-authored-by: Alex Waygood <[email protected]>
1 parent 976c37a commit 3f6c65e

File tree

2 files changed

+82
-12
lines changed

2 files changed

+82
-12
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
# Consolidating narrowed types after if statement
2+
3+
## After if-else statements, narrowing has no effect if the variable is not mutated in any branch
4+
5+
```py
6+
def optional_int() -> int | None: ...
7+
8+
x = optional_int()
9+
10+
if x is None:
11+
pass
12+
else:
13+
pass
14+
15+
reveal_type(x) # revealed: int | None
16+
```
17+
18+
## Narrowing can have a persistent effect if the variable is mutated in one branch
19+
20+
```py
21+
def optional_int() -> int | None: ...
22+
23+
x = optional_int()
24+
25+
if x is None:
26+
x = 10
27+
else:
28+
pass
29+
30+
reveal_type(x) # revealed: int
31+
```
32+
33+
## An if statement without an explicit `else` branch is equivalent to one with a no-op `else` branch
34+
35+
```py
36+
def optional_int() -> int | None: ...
37+
38+
x = optional_int()
39+
y = optional_int()
40+
41+
if x is None:
42+
x = 0
43+
44+
if y is None:
45+
pass
46+
47+
reveal_type(x) # revealed: int
48+
reveal_type(y) # revealed: int | None
49+
```
50+
51+
## An if-elif without an explicit else branch is equivalent to one with an empty else branch
52+
53+
```py
54+
def optional_int() -> int | None: ...
55+
56+
x = optional_int()
57+
58+
if x is None:
59+
x = 0
60+
elif x > 50:
61+
x = 50
62+
63+
reveal_type(x) # revealed: int
64+
```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

+18-12
Original file line numberDiff line numberDiff line change
@@ -769,7 +769,22 @@ where
769769
let mut constraints = vec![constraint];
770770
self.visit_body(&node.body);
771771
let mut post_clauses: Vec<FlowSnapshot> = vec![];
772-
for clause in &node.elif_else_clauses {
772+
let elif_else_clauses = node
773+
.elif_else_clauses
774+
.iter()
775+
.map(|clause| (clause.test.as_ref(), clause.body.as_slice()));
776+
let has_else = node
777+
.elif_else_clauses
778+
.last()
779+
.is_some_and(|clause| clause.test.is_none());
780+
let elif_else_clauses = elif_else_clauses.chain(if has_else {
781+
// if there's an `else` clause already, we don't need to add another
782+
None
783+
} else {
784+
// if there's no `else` branch, we should add a no-op `else` branch
785+
Some((None, Default::default()))
786+
});
787+
for (clause_test, clause_body) in elif_else_clauses {
773788
// snapshot after every block except the last; the last one will just become
774789
// the state that we merge the other snapshots into
775790
post_clauses.push(self.flow_snapshot());
@@ -779,24 +794,15 @@ where
779794
for constraint in &constraints {
780795
self.record_negated_constraint(*constraint);
781796
}
782-
if let Some(elif_test) = &clause.test {
797+
if let Some(elif_test) = clause_test {
783798
self.visit_expr(elif_test);
784799
constraints.push(self.record_expression_constraint(elif_test));
785800
}
786-
self.visit_body(&clause.body);
801+
self.visit_body(clause_body);
787802
}
788803
for post_clause_state in post_clauses {
789804
self.flow_merge(post_clause_state);
790805
}
791-
let has_else = node
792-
.elif_else_clauses
793-
.last()
794-
.is_some_and(|clause| clause.test.is_none());
795-
if !has_else {
796-
// if there's no else clause, then it's possible we took none of the branches,
797-
// and the pre_if state can reach here
798-
self.flow_merge(pre_if);
799-
}
800806
}
801807
ast::Stmt::While(ast::StmtWhile {
802808
test,

0 commit comments

Comments
 (0)