fix narrowing on dict (and typed dict) .get

pyrefly/lib/alt/narrow.rs

@@ -26,6 +26,8 @@ use crate::alt::answers_solver::AnswersSolver;
 use crate::alt::callable::CallArg;
 use crate::alt::callable::CallKeyword;
 use crate::binding::narrow::AtomicNarrowOp;
+use crate::binding::narrow::FacetOrigin;
+use crate::binding::narrow::FacetSubject;
 use crate::binding::narrow::NarrowOp;
 use crate::error::collector::ErrorCollector;
 use crate::error::style::ErrorStyle;
@@ -809,15 +811,17 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         match op {
             NarrowOp::Atomic(subject, AtomicNarrowOp::HasAttr(attr)) => {
                 let base_ty = match subject {
-                    Some(facet_chain) => self.get_facet_chain_type(type_info, facet_chain, range),
+                    Some(facet_subject) => {
+                        self.get_facet_chain_type(type_info, &facet_subject.chain, range)
+                    }
                     None => type_info.ty().clone(),
                 };
                 // We only narrow the attribute to `Any` if the attribute does not exist
                 if !self.has_attr(&base_ty, attr) {
                     let attr_facet = FacetKind::Attribute(attr.clone());
                     let facets = match subject {
-                        Some(chain) => {
-                            let mut new_facets = chain.facets().clone();
+                        Some(facet_subject) => {
+                            let mut new_facets = facet_subject.chain.facets().clone();
                             new_facets.push(attr_facet);
                             new_facets
                         }
@@ -839,15 +843,17 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                     return type_info.clone();
                 }
                 let base_ty = match subject {
-                    Some(facet_chain) => self.get_facet_chain_type(type_info, facet_chain, range),
+                    Some(facet_subject) => {
+                        self.get_facet_chain_type(type_info, &facet_subject.chain, range)
+                    }
                     None => type_info.ty().clone(),
                 };
                 let attr_ty =
                     self.attr_infer_for_type(&base_ty, attr, range, &suppress_errors, None);
                 let attr_facet = FacetKind::Attribute(attr.clone());
                 let facets = match subject {
-                    Some(chain) => {
-                        let mut new_facets = chain.facets().clone();
+                    Some(facet_subject) => {
+                        let mut new_facets = facet_subject.chain.facets().clone();
                         new_facets.push(attr_facet);
                         new_facets
                     }
@@ -872,16 +878,21 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 let ty = self.atomic_narrow(type_info.ty(), op, range, errors);
                 type_info.clone().with_ty(ty)
             }
-            NarrowOp::Atomic(Some(facet_chain), op) => {
+            NarrowOp::Atomic(Some(facet_subject), op) => {
+                if facet_subject.origin == FacetOrigin::DictGet
+                    && !self.supports_dict_get_subject(type_info, &facet_subject, range)
+                {
+                    return type_info.clone();
+                }
                 let ty = self.atomic_narrow(
-                    &self.get_facet_chain_type(type_info, facet_chain, range),
+                    &self.get_facet_chain_type(type_info, &facet_subject.chain, range),
                     op,
                     range,
                     errors,
                 );
-                let mut narrowed = type_info.with_narrow(facet_chain.facets(), ty);
+                let mut narrowed = type_info.with_narrow(facet_subject.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() {
+                if let Some((last, prefix)) = facet_subject.chain.facets().split_last() {
                     match Vec1::try_from(prefix) {
                         Ok(prefix_facets) => {
                             let prefix_chain = FacetChain::new(prefix_facets);
@@ -927,4 +938,50 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             ),
         }
     }
+
+    fn supports_dict_get_subject(
+        &self,
+        type_info: &TypeInfo,
+        subject: &FacetSubject,
+        range: TextRange,
+    ) -> bool {
+        let base_ty = if subject.chain.facets().len() == 1 {
+            type_info.ty().clone()
+        } else {
+            let prefix: Vec<_> = subject
+                .chain
+                .facets()
+                .iter()
+                .take(subject.chain.facets().len() - 1)
+                .cloned()
+                .collect();
+            match Vec1::try_from_vec(prefix) {
+                Ok(vec1) => {
+                    let prefix_chain = FacetChain::new(vec1);
+                    self.get_facet_chain_type(type_info, &prefix_chain, range)
+                }
+                Err(_) => return false,
+            }
+        };
+        self.is_dict_like_type(&base_ty)
+    }
+
+    fn is_dict_like_type(&self, ty: &Type) -> bool {
+        match ty {
+            Type::ClassType(cls) => cls.is_builtin("dict"),
+            Type::TypedDict(_) | Type::PartialTypedDict(_) => true,
+            Type::Union(types) => types.iter().all(|t| self.is_dict_like_type(t)),
+            Type::TypeAlias(alias) => self.is_dict_like_type(&alias.as_value(self.stdlib)),
+            Type::Var(var) => {
+                let Some(_guard) = self.recurse(*var) else {
+                    return false;
+                };
+                let forced = self.solver().force_var(*var);
+                self.is_dict_like_type(&forced)
+            }
+            Type::Forall(forall) => self.is_dict_like_type(&forall.body.clone().as_type()),
+            Type::Type(box inner) => self.is_dict_like_type(inner),
+            _ => false,
+        }
+    }
 }

pyrefly/lib/binding/narrow.rs

@@ -44,7 +44,7 @@ use crate::types::facet::FacetKind;
 use crate::types::types::Type;
 
 assert_words!(AtomicNarrowOp, 11);
-assert_words!(NarrowOp, 12);
+assert_words!(NarrowOp, 13);
 
 #[derive(Clone, Debug)]
 pub enum AtomicNarrowOp {
@@ -92,7 +92,7 @@ pub enum AtomicNarrowOp {
 
 #[derive(Clone, Debug)]
 pub enum NarrowOp {
-    Atomic(Option<FacetChain>, AtomicNarrowOp),
+    Atomic(Option<FacetSubject>, AtomicNarrowOp),
     And(Vec<NarrowOp>),
     Or(Vec<NarrowOp>),
 }
@@ -176,7 +176,7 @@ impl DisplayWith<ModuleInfo> for NarrowOp {
         match self {
             Self::Atomic(prop, op) => match prop {
                 None => write!(f, "{}", op.display_with(ctx)),
-                Some(prop) => write!(f, "[{prop}] {}", op.display_with(ctx)),
+                Some(prop) => write!(f, "[{}] {}", prop.chain, op.display_with(ctx)),
             },
             Self::And(ops) => {
                 write!(
@@ -234,19 +234,43 @@ impl AtomicNarrowOp {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum FacetOrigin {
+    Direct,
+    DictGet,
+}
+
+#[derive(Clone, Debug)]
+pub struct FacetSubject {
+    pub chain: FacetChain,
+    pub origin: FacetOrigin,
+}
+
 #[derive(Clone, Debug)]
 pub enum NarrowingSubject {
     Name(Name),
-    Facets(Name, FacetChain),
+    Facets(Name, FacetSubject),
 }
 
 impl NarrowingSubject {
     pub fn with_facet(&self, prop: FacetKind) -> Self {
         match self {
-            Self::Name(name) => Self::Facets(name.clone(), FacetChain::new(Vec1::new(prop))),
-            Self::Facets(name, props) => {
-                let props = Vec1::from_vec_push(props.facets().to_vec(), prop);
-                Self::Facets(name.clone(), FacetChain::new(props))
+            Self::Name(name) => Self::Facets(
+                name.clone(),
+                FacetSubject {
+                    chain: FacetChain::new(Vec1::new(prop)),
+                    origin: FacetOrigin::Direct,
+                },
+            ),
+            Self::Facets(name, facets) => {
+                let props = Vec1::from_vec_push(facets.chain.facets().to_vec(), prop);
+                Self::Facets(
+                    name.clone(),
+                    FacetSubject {
+                        chain: FacetChain::new(props),
+                        origin: facets.origin,
+                    },
+                )
             }
         }
     }
@@ -372,7 +396,7 @@ impl NarrowOps {
         for subject in expr_to_subjects(left) {
             let (name, prop) = match subject {
                 NarrowingSubject::Name(name) => (name, None),
-                NarrowingSubject::Facets(name, prop) => (name, Some(prop)),
+                NarrowingSubject::Facets(name, facets) => (name, Some(facets)),
             };
             if let Some((existing, _)) = narrow_ops.0.get_mut(&name) {
                 existing.and(NarrowOp::Atomic(prop, op.clone()));
@@ -393,7 +417,7 @@ impl NarrowOps {
         let mut narrow_ops = Self::new();
         let (name, prop) = match subject {
             NarrowingSubject::Name(name) => (name, None),
-            NarrowingSubject::Facets(name, prop) => (name, Some(prop)),
+            NarrowingSubject::Facets(name, facets) => (name, Some(facets)),
         };
         if let Some((existing, _)) = narrow_ops.0.get_mut(&name) {
             existing.and(NarrowOp::Atomic(prop, op.clone()));
@@ -771,15 +795,58 @@ pub fn identifier_and_chain_prefix_for_expr(expr: &Expr) -> Option<(Identifier,
 }
 
 fn subject_for_expr(expr: &Expr) -> Option<NarrowingSubject> {
-    identifier_and_chain_for_expr(expr)
-        .map(|(identifier, attr)| NarrowingSubject::Facets(identifier.id, attr))
+    if let Some((identifier, attr)) = identifier_and_chain_for_expr(expr) {
+        return Some(NarrowingSubject::Facets(
+            identifier.id,
+            FacetSubject {
+                chain: attr,
+                origin: FacetOrigin::Direct,
+            },
+        ));
+    }
+
+    if let Expr::Call(ExprCall {
+        func, arguments, ..
+    }) = expr
+        && arguments.keywords.is_empty()
+        && let Some(first_arg) = arguments.args.first()
+        && let Expr::Attribute(attr) = &**func
+        && attr.attr.id.as_str() == "get"
+        && let Expr::StringLiteral(ExprStringLiteral { value, .. }) = first_arg
+    {
+        let key = value.to_string();
+        if let Some((identifier, facets)) = identifier_and_chain_for_expr(&attr.value) {
+            let props = Vec1::from_vec_push(facets.facets().to_vec(), FacetKind::Key(key.clone()));
+            return Some(NarrowingSubject::Facets(
+                identifier.id,
+                FacetSubject {
+                    chain: FacetChain::new(props),
+                    origin: FacetOrigin::DictGet,
+                },
+            ));
+        } else if let Expr::Name(name) = &*attr.value {
+            return Some(NarrowingSubject::Facets(
+                name.id.clone(),
+                FacetSubject {
+                    chain: FacetChain::new(Vec1::new(FacetKind::Key(key))),
+                    origin: FacetOrigin::DictGet,
+                },
+            ));
+        }
+    }
+
+    None
 }
 
 pub fn expr_to_subjects(expr: &Expr) -> Vec<NarrowingSubject> {
     fn f(expr: &Expr, res: &mut Vec<NarrowingSubject>) {
         match expr {
             Expr::Name(name) => res.push(NarrowingSubject::Name(name.id.clone())),
-            Expr::Attribute(_) | Expr::Subscript(_) => res.extend(subject_for_expr(expr)),
+            Expr::Attribute(_) | Expr::Subscript(_) | Expr::Call(_) => {
+                if let Some(subject) = subject_for_expr(expr) {
+                    res.push(subject);
+                }
+            }
             Expr::Named(ExprNamed { target, value, .. }) => {
                 f(target, res);
                 f(value, res);

pyrefly/lib/test/subscript_narrow.rs

@@ -183,10 +183,72 @@ class ErrorContext:
             self.system_context = {"status": "active"}
 
         assert_type(self.system_context, dict[str, Any])
-            
+
         self.system_context["timestamp"] = "2024-01-01"
 
         assert_type(self.system_context, dict[str, Any])
         assert_type(self.system_context["timestamp"], Literal["2024-01-01"])
 "#,
 );
+
+testcase!(
+    bug = "https://github.com/facebook/pyrefly/issues/238",
+    test_dict_get_literal_key_narrow,
+    r#"
+from typing import assert_type
+
+def narrow_with_explicit_none(data: dict[str, int]) -> None:
+    value = data.get("foo")
+    if value is not None:
+        assert_type(value, int)
+        assert_type(data["foo"], int)
+    else:
+        assert_type(value, None)
+
+def narrow_with_truthy_check(data: dict[str, int]) -> None:
+    if data.get("bar"):
+        assert_type(data["bar"], int)
+    else:
+        fallback = data.get("bar")
+        assert_type(fallback, int | None)
+"#,
+);
+
+testcase!(
+    bug = "https://github.com/facebook/pyrefly/issues/238",
+    test_typeddict_get_literal_key_narrow,
+    r#"
+from typing import TypedDict, assert_type
+
+class TD(TypedDict, total=False):
+    foo: int
+
+def use(td: TD) -> None:
+    value = td.get("foo")
+    if value is not None:
+        assert_type(value, int)
+        assert_type(td["foo"], int)
+    else:
+        assert_type(value, None)
+"#,
+);
+
+testcase!(
+    bug = "https://github.com/facebook/pyrefly/issues/238",
+    test_non_dict_get_does_not_narrow,
+    r#"
+from typing import assert_type
+
+class NotDict:
+    def get(self, key: str) -> int | None: ...
+    def __getitem__(self, key: str) -> int | None: ...
+
+def use(mapping: NotDict) -> None:
+    if mapping.get("foo") is not None:
+        assert_type(mapping.get("foo"), int | None)
+        assert_type(mapping["foo"], int | None)
+    else:
+        assert_type(mapping.get("foo"), int | None)
+        assert_type(mapping["foo"], int | None)
+"#,
+);