From 1f5438b30213322d39a5ac5beec6e31ffe10065c Mon Sep 17 00:00:00 2001
From: Rex Ledesma <jledesma@nvidia.com>
Date: Sat, 17 Oct 2026 11:12:54 -0400
Subject: [PATCH] fix(narrowing): preserve reassigned facet types across
 branches

Attribute and subscript assignments can widen facet types without rebinding their root name. Treat widened facet types as reassignments when merging branches so the merge preserves the new values instead of reverting to the earlier type.
---
 crates/pyrefly_types/src/type_info.rs | 33 ++++++++++++++++-----------
 pyrefly/lib/test/attribute_narrow.rs  | 31 +++++++++++++++++++++++++
 pyrefly/lib/test/subscript_narrow.rs  | 16 +++++++++++++
 3 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/crates/pyrefly_types/src/type_info.rs b/crates/pyrefly_types/src/type_info.rs
index 83ff5ee097..76dd4c1e8b 100644
--- a/crates/pyrefly_types/src/type_info.rs
+++ b/crates/pyrefly_types/src/type_info.rs
@@ -850,19 +850,26 @@ impl NarrowedFacet {
             }
         }
         let ty = ty_branches.map(|tys| {
-            join_types(
-                tys,
-                union_types,
-                is_subset_eq,
-                join_style.flat_map(|base_facet| {
-                    base_facet.as_ref().and_then(|facet| match facet {
-                        NarrowedFacet::WithRoot(ty, _) | NarrowedFacet::Leaf(ty) => {
-                            Some(ty.clone())
-                        }
-                        NarrowedFacet::WithoutRoot { .. } => None,
-                    })
-                }),
-            )
+            let join_style = join_style.flat_map(|base_facet| {
+                base_facet.as_ref().and_then(|facet| match facet {
+                    NarrowedFacet::WithRoot(ty, _) | NarrowedFacet::Leaf(ty) => Some(ty.clone()),
+                    NarrowedFacet::WithoutRoot { .. } => None,
+                })
+            });
+            // Assigning a facet does not rebind its root name. A merge of root narrows
+            // can therefore contain facet reassignments that widen the base type.
+            // Such facets must not simplify back to the type before the branch.
+            let join_style = match join_style {
+                JoinStyle::NarrowOf(base_ty)
+                    if tys
+                        .iter()
+                        .any(|ty| ty != &base_ty && !is_subset_eq(ty, &base_ty)) =>
+                {
+                    JoinStyle::ReassignmentOf(base_ty)
+                }
+                style => style,
+            };
+            join_types(tys, union_types, is_subset_eq, join_style)
         });
         let facets = facets_branches.and_then(|branches| {
             NarrowedFacets::join(
diff --git a/pyrefly/lib/test/attribute_narrow.rs b/pyrefly/lib/test/attribute_narrow.rs
index f7295f2516..bc77648e82 100644
--- a/pyrefly/lib/test/attribute_narrow.rs
+++ b/pyrefly/lib/test/attribute_narrow.rs
@@ -7,6 +7,37 @@
 
 use crate::testcase;
 
+testcase!(
+    test_conditional_attribute_reassignment_preserves_union,
+    r#"
+from typing import Literal, assert_type
+class Foo:
+    x: bool = True
+class Bar:
+    def __init__(self, flag: bool) -> None:
+        self.x: Foo | None = None
+        if flag:
+            self.x = Foo()
+        assert_type(self.x, Foo | None)
+        assert self.x is not None
+        assert_type(self.x, Foo)
+        self.x.x = False
+class Baz:
+    x: int | str
+def f(baz: Baz, flag: bool) -> None:
+    baz.x = 0
+    if flag:
+        baz.x = "a"
+    assert_type(baz.x, Literal[0, "a"])
+def g(bar: Bar, flag: bool) -> None:
+    bar.x = Foo()
+    bar.x.x = True
+    if flag:
+        bar.x.x = False
+    assert_type(bar.x.x, bool)
+"#,
+);
+
 testcase!(
     test_one_level_simple,
     r#"
diff --git a/pyrefly/lib/test/subscript_narrow.rs b/pyrefly/lib/test/subscript_narrow.rs
index 866f714db4..54207cd24a 100644
--- a/pyrefly/lib/test/subscript_narrow.rs
+++ b/pyrefly/lib/test/subscript_narrow.rs
@@ -8,6 +8,22 @@
 use crate::test::util::TestEnv;
 use crate::testcase;
 
+testcase!(
+    test_conditional_subscript_reassignment_preserves_union,
+    r#"
+from typing import Literal, TypedDict, assert_type
+class C(TypedDict):
+    x: int | None
+def f(c: C, flag: bool) -> None:
+    c["x"] = None
+    if flag:
+        c["x"] = 1
+    assert_type(c["x"], Literal[1] | None)
+    assert c["x"] is not None
+    assert_type(c["x"], Literal[1])
+"#,
+);
+
 testcase!(
     test_index_narrow,
     r#"