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