From f98eebdbab73d82a9db53f16bf6f34577ebd345e Mon Sep 17 00:00:00 2001
From: cake-monotone <cake.monotone@gmail.com>
Date: Tue, 26 Nov 2024 03:36:37 +0900
Subject: [PATCH] [red-knot] Fix Leaking Narrowing Constraint in `ast::ExprIf`
 (#14590)

## Summary

Closes #14588


```py
x: Literal[42, "hello"] = 42 if bool_instance() else "hello"
reveal_type(x)  # revealed: Literal[42] | Literal["hello"]

_ = ... if isinstance(x, str) else ...

# The `isinstance` test incorrectly narrows the type of `x`.
# As a result, `x` is revealed as Literal["hello"], but it should remain Literal[42, "hello"].
reveal_type(x)  # revealed: Literal["hello"]
```

## Test Plan
mdtest included!

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 .../resources/mdtest/expression/if.md         | 19 +++++++++++++++++++
 .../src/semantic_index/builder.rs             |  2 +-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/expression/if.md b/crates/red_knot_python_semantic/resources/mdtest/expression/if.md
index ec687f798f..68d5ef444f 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/expression/if.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/expression/if.md
@@ -22,3 +22,22 @@ reveal_type(1 if None else 2)  # revealed: Literal[2]
 reveal_type(1 if "" else 2)  # revealed: Literal[2]
 reveal_type(1 if 0 else 2)  # revealed: Literal[2]
 ```
+
+## Leaked Narrowing Constraint
+
+(issue #14588)
+
+The test inside an if expression should not affect code outside of the block.
+
+```py
+def bool_instance() -> bool:
+    return True
+
+x: Literal[42, "hello"] = 42 if bool_instance() else "hello"
+
+reveal_type(x)  # revealed: Literal[42] | Literal["hello"]
+
+_ = ... if isinstance(x, str) else ...
+
+reveal_type(x)  # revealed: Literal[42] | Literal["hello"]
+```
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 dfca6328db..1c4459a777 100644
--- a/crates/red_knot_python_semantic/src/semantic_index/builder.rs
+++ b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
@@ -1189,8 +1189,8 @@ where
                 // AST inspection, so we can't simplify here, need to record test expression for
                 // later checking)
                 self.visit_expr(test);
-                let constraint = self.record_expression_constraint(test);
                 let pre_if = self.flow_snapshot();
+                let constraint = self.record_expression_constraint(test);
                 self.visit_expr(body);
                 let post_body = self.flow_snapshot();
                 self.flow_restore(pre_if);