support discriminated unions in narrowing

Summary:
When we have a check like `x.y == 1` or `x["key"] = "value"` or `x[0] is None`, we previously only narrowed the appropriate facet for `x`. This diff makes us also narrow the type of `x`.

This is useful for modeling a form of tagged unions in user code, where we have a union of classes/typeddicts/tuples that's differentiated by a facet that's typed as some literal.

This is a feature that's supported for Pyright, and should be useful for Pydantic users as well.

see: #650

closes #418

Reviewed By: stroxler

Differential Revision: D80464317

fbshipit-source-id: ec583bbb910c53fecff4384b292cf8adc9ea93a9

Author: yangdanny97

pyrefly/lib/alt/narrow.rs

@@ -20,6 +20,7 @@ use ruff_python_ast::StringLiteralValue;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
+use vec1::Vec1;
 
 use crate::alt::answers::LookupAnswer;
 use crate::alt::answers_solver::AnswersSolver;
@@ -240,6 +241,102 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         })
     }
 
+    // Try to narrow a type based on the type of its facet.
+    // For example, if we have a `x.y == 0` check and `x` is some union,
+    // we can eliminate cases from the union where `x.y` is some other
+    // literal.
+    pub fn atomic_narrow_for_facet(
+        &self,
+        base: &Type,
+        facet: &FacetKind,
+        op: &AtomicNarrowOp,
+        range: TextRange,
+        errors: &ErrorCollector,
+    ) -> Option<Type> {
+        match op {
+            AtomicNarrowOp::Is(v) => {
+                let right = self.expr_infer(v, errors);
+                Some(self.distribute_over_union(base, |t| {
+                    let base_info = TypeInfo::of_ty(t.clone());
+                    let facet_ty = self.get_facet_chain_type(
+                        &base_info,
+                        &FacetChain::new(Vec1::new(facet.clone())),
+                        range,
+                    );
+                    match right {
+                        Type::None | Type::Literal(Lit::Bool(_)) | Type::Literal(Lit::Enum(_)) => {
+                            if self.is_subset_eq(&right, &facet_ty) {
+                                t.clone()
+                            } else {
+                                Type::never()
+                            }
+                        }
+                        _ => t.clone(),
+                    }
+                }))
+            }
+            AtomicNarrowOp::IsNot(v) => {
+                let right = self.expr_infer(v, errors);
+                Some(self.distribute_over_union(base, |t| {
+                    let base_info = TypeInfo::of_ty(t.clone());
+                    let facet_ty = self.get_facet_chain_type(
+                        &base_info,
+                        &FacetChain::new(Vec1::new(facet.clone())),
+                        range,
+                    );
+                    match (&facet_ty, &right) {
+                        (
+                            Type::None | Type::Literal(Lit::Bool(_)) | Type::Literal(Lit::Enum(_)),
+                            Type::None | Type::Literal(Lit::Bool(_)) | Type::Literal(Lit::Enum(_)),
+                        ) if right == facet_ty => Type::never(),
+                        _ => t.clone(),
+                    }
+                }))
+            }
+            AtomicNarrowOp::Eq(v) => {
+                let right = self.expr_infer(v, errors);
+                Some(self.distribute_over_union(base, |t| {
+                    let base_info = TypeInfo::of_ty(t.clone());
+                    let facet_ty = self.get_facet_chain_type(
+                        &base_info,
+                        &FacetChain::new(Vec1::new(facet.clone())),
+                        range,
+                    );
+                    match right {
+                        Type::None | Type::Literal(_) => {
+                            if self.is_subset_eq(&right, &facet_ty) {
+                                t.clone()
+                            } else {
+                                Type::never()
+                            }
+                        }
+                        _ => t.clone(),
+                    }
+                }))
+            }
+            AtomicNarrowOp::NotEq(v) => {
+                let right = self.expr_infer(v, errors);
+                Some(self.distribute_over_union(base, |t| {
+                    let base_info = TypeInfo::of_ty(t.clone());
+                    let facet_ty = self.get_facet_chain_type(
+                        &base_info,
+                        &FacetChain::new(Vec1::new(facet.clone())),
+                        range,
+                    );
+                    match (&facet_ty, &right) {
+                        (Type::None | Type::Literal(_), Type::None | Type::Literal(_))
+                            if right == facet_ty =>
+                        {
+                            Type::never()
+                        }
+                        _ => t.clone(),
+                    }
+                }))
+            }
+            _ => None,
+        }
+    }
+
     pub fn atomic_narrow(
         &self,
         ty: &Type,
@@ -709,7 +806,33 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                     range,
                     errors,
                 );
-                type_info.with_narrow(facet_chain.facets(), ty)
+                let mut narrowed = type_info.with_narrow(facet_chain.facets(), ty);
+                // For certain types of narrows, we can also narrow the parent of the current subject
+                if let Some((last, prefix)) = facet_chain.facets().split_last() {
+                    match Vec1::try_from(prefix) {
+                        Ok(prefix_facets) => {
+                            let prefix_chain = FacetChain::new(prefix_facets);
+                            let base_ty =
+                                self.get_facet_chain_type(type_info, &prefix_chain, range);
+                            if let Some(narrowed_ty) =
+                                self.atomic_narrow_for_facet(&base_ty, last, op, range, errors)
+                                && narrowed_ty != base_ty
+                            {
+                                narrowed = narrowed.with_narrow(prefix_chain.facets(), narrowed_ty);
+                            }
+                        }
+                        _ => {
+                            let base_ty = type_info.ty();
+                            if let Some(narrowed_ty) =
+                                self.atomic_narrow_for_facet(base_ty, last, op, range, errors)
+                                && narrowed_ty != *base_ty
+                            {
+                                narrowed = narrowed.clone().with_ty(narrowed_ty);
+                            }
+                        }
+                    };
+                }
+                narrowed
             }
             NarrowOp::And(ops) => {
                 let mut ops_iter = ops.iter();

pyrefly/lib/test/narrow.rs

@@ -1643,3 +1643,77 @@ class Test:
                 pass
     "#,
 );
+
+testcase!(
+    test_discriminated_union_key,
+    r#"
+from typing import TypedDict, assert_type, Literal
+
+class UserDict(TypedDict):
+    kind: Literal["user"]
+    is_admin: Literal[False]
+
+class AdminDict(TypedDict):
+    kind: Literal["admin"]
+    is_admin: Literal[True]
+
+def test(x: UserDict | AdminDict):
+    if x["kind"] == "user":
+        assert_type(x, UserDict)
+    else:
+        assert_type(x, AdminDict)
+    if x["is_admin"] is True:
+        assert_type(x, AdminDict)
+    else:
+        assert_type(x, UserDict)
+    "#,
+);
+
+testcase!(
+    test_discriminated_union_attr,
+    r#"
+from typing import assert_type, Literal
+
+class User:
+    kind: Literal["user"]
+    is_admin: Literal[False]
+
+class Admin:
+    kind: Literal["admin"]
+    is_admin: Literal[True]
+
+def test(x: User | Admin):
+    if x.kind == "user":
+        assert_type(x, User)
+    elif x.kind == "admin":
+        assert_type(x, Admin)
+
+    if x.is_admin is True:
+        assert_type(x, Admin)
+    else:
+        assert_type(x, User)
+    "#,
+);
+
+testcase!(
+    test_discriminated_union_index,
+    r#"
+from typing import assert_type, Literal
+
+def test(x: tuple[Literal["user"], Literal[None]] | tuple[Literal["admin"], int]):
+    if x[0] == "user":
+        assert_type(x, tuple[Literal["user"], Literal[None]])
+    else:
+        assert_type(x, tuple[Literal["admin"], int])
+
+    if x[1] is None:
+        assert_type(x, tuple[Literal["user"], Literal[None]])
+    else:
+        assert_type(x, tuple[Literal["admin"], int])
+
+    if x[1] is not None:
+        assert_type(x, tuple[Literal["admin"], int])
+    else:
+        assert_type(x, tuple[Literal["user"], Literal[None]])  
+    "#,
+);