From 3f6c65e78c10579d21e1513aa59ea6dc0248f008 Mon Sep 17 00:00:00 2001
From: Samodya Abeysiriwardane <379594+sransara@users.noreply.github.com>
Date: Thu, 28 Nov 2024 08:23:55 -0600
Subject: [PATCH] [red-knot] Fix merged type after if-else without explicit
 else branch (#14621)

## Summary

Closes: https://github.com/astral-sh/ruff/issues/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 <carl@astral.sh>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 .../mdtest/narrow/post_if_statement.md        | 64 +++++++++++++++++++
 .../src/semantic_index/builder.rs             | 30 +++++----
 2 files changed, 82 insertions(+), 12 deletions(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/narrow/post_if_statement.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/post_if_statement.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/post_if_statement.md
new file mode 100644
index 00000000000000..267e820d595da4
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/post_if_statement.md
@@ -0,0 +1,64 @@
+# Consolidating narrowed types after if statement
+
+## After if-else statements, narrowing has no effect if the variable is not mutated in any branch
+
+```py
+def optional_int() -> int | None: ...
+
+x = optional_int()
+
+if x is None:
+    pass
+else:
+    pass
+
+reveal_type(x)  # revealed: int | None
+```
+
+## Narrowing can have a persistent effect if the variable is mutated in one branch
+
+```py
+def optional_int() -> int | None: ...
+
+x = optional_int()
+
+if x is None:
+    x = 10
+else:
+    pass
+
+reveal_type(x)  # revealed: int
+```
+
+## An if statement without an explicit `else` branch is equivalent to one with a no-op `else` branch
+
+```py
+def optional_int() -> int | None: ...
+
+x = optional_int()
+y = optional_int()
+
+if x is None:
+    x = 0
+
+if y is None:
+    pass
+
+reveal_type(x)  # revealed: int
+reveal_type(y)  # revealed: int | None
+```
+
+## An if-elif without an explicit else branch is equivalent to one with an empty else branch
+
+```py
+def optional_int() -> int | None: ...
+
+x = optional_int()
+
+if x is None:
+    x = 0
+elif x > 50:
+    x = 50
+
+reveal_type(x)  # revealed: int
+```
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 1c4459a777f680..68258c9759153d 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
@@ -769,7 +769,22 @@ where
                 let mut constraints = vec![constraint];
                 self.visit_body(&node.body);
                 let mut post_clauses: Vec<FlowSnapshot> = vec![];
-                for clause in &node.elif_else_clauses {
+                let elif_else_clauses = node
+                    .elif_else_clauses
+                    .iter()
+                    .map(|clause| (clause.test.as_ref(), clause.body.as_slice()));
+                let has_else = node
+                    .elif_else_clauses
+                    .last()
+                    .is_some_and(|clause| clause.test.is_none());
+                let elif_else_clauses = elif_else_clauses.chain(if has_else {
+                    // if there's an `else` clause already, we don't need to add another
+                    None
+                } else {
+                    // if there's no `else` branch, we should add a no-op `else` branch
+                    Some((None, Default::default()))
+                });
+                for (clause_test, clause_body) in elif_else_clauses {
                     // snapshot after every block except the last; the last one will just become
                     // the state that we merge the other snapshots into
                     post_clauses.push(self.flow_snapshot());
@@ -779,24 +794,15 @@ where
                     for constraint in &constraints {
                         self.record_negated_constraint(*constraint);
                     }
-                    if let Some(elif_test) = &clause.test {
+                    if let Some(elif_test) = clause_test {
                         self.visit_expr(elif_test);
                         constraints.push(self.record_expression_constraint(elif_test));
                     }
-                    self.visit_body(&clause.body);
+                    self.visit_body(clause_body);
                 }
                 for post_clause_state in post_clauses {
                     self.flow_merge(post_clause_state);
                 }
-                let has_else = node
-                    .elif_else_clauses
-                    .last()
-                    .is_some_and(|clause| clause.test.is_none());
-                if !has_else {
-                    // if there's no else clause, then it's possible we took none of the branches,
-                    // and the pre_if state can reach here
-                    self.flow_merge(pre_if);
-                }
             }
             ast::Stmt::While(ast::StmtWhile {
                 test,