[red-knot] Fix merged type after if-else without explicit else branch

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
+```

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
 use except_handlers::TryNodeContextStackManager;
+use itertools::Itertools;
 use rustc_hash::FxHashMap;
 
 use ruff_db::files::File;
@@ -768,8 +769,22 @@ where
                 let constraint = self.record_expression_constraint(&node.test);
                 let mut constraints = vec![constraint];
                 self.visit_body(&node.body);
+                let mut elif_else_clauses = node
+                    .elif_else_clauses
+                    .iter()
+                    .map(|clause| (&clause.test, &clause.body[..]))
+                    .collect_vec();
+                let has_else = elif_else_clauses
+                    .last()
+                    .is_some_and(|clause| clause.0.is_none());
+                if !has_else {
+                    // an if-elif statement without an explicit `else` branch is equivalent to one with a no-op `else` branch
+                    elif_else_clauses.push((&None, &[]));
+                }
                 let mut post_clauses: Vec<FlowSnapshot> = vec![];
-                for clause in &node.elif_else_clauses {
+                for clause in elif_else_clauses {
+                    let clause_test = clause.0;
+                    let clause_body = clause.1;
                     // 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,