fix(narrowing): preserve reassigned facet types across branches

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(

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#"

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#"