fix narrowing on dict (and typed dict) .get (#1390)

Summary:
fix #238

Taught the narrowing logic to treat dict.get("literal") calls as key facets, so truthiness and is not None checks now refine the corresponding dictionary entry

Added regression coverage for both explicit is not None and plain truthiness checks on literal-key .get calls to confirm the dictionary entry narrows as expected

Pull Request resolved: #1390

Reviewed By: stroxler

Differential Revision: D85959099

Pulled By: yangdanny97

fbshipit-source-id: a7d1b4f5983f1f21486a837f80e0001e85128adf

Author: asukaminato0721

crates/pyrefly_types/src/types.rs

@@ -888,6 +888,10 @@ impl Type {
         matches!(self, Type::Any(_))
     }
 
+    pub fn is_typed_dict(&self) -> bool {
+        matches!(self, Type::TypedDict(_) | Type::PartialTypedDict(_))
+    }
+
     pub fn is_error(&self) -> bool {
         matches!(self, Type::Any(AnyStyle::Error))
     }

pyrefly/lib/alt/expr.rs

@@ -38,6 +38,7 @@ use ruff_python_ast::ExprTuple;
 use ruff_python_ast::Identifier;
 use ruff_python_ast::Keyword;
 use ruff_python_ast::Number;
+use ruff_python_ast::StringLiteralValue;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
@@ -542,6 +543,14 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 let callee_ty = self.expr_infer(&x.func, errors);
                 if let Some(d) = self.call_to_dict(&callee_ty, &x.arguments) {
                     self.dict_infer(&d, hint, x.range, errors)
+                } else if let Some((obj_ty, key)) =
+                    self.is_dict_get_with_literal(&x.func, &x.arguments, errors)
+                {
+                    obj_ty
+                        .at_facet(&FacetKind::Key(key.to_string()), || {
+                            self.expr_call_infer(x, callee_ty.clone(), hint, errors)
+                        })
+                        .into_ty()
                 } else {
                     self.expr_call_infer(x, callee_ty, hint, errors)
                 }
@@ -971,6 +980,48 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         }))
     }
 
+    // Is this a call to `dict.get` with a single string literal argument
+    fn is_dict_get_with_literal(
+        &self,
+        func: &Expr,
+        args: &Arguments,
+        errors: &ErrorCollector,
+    ) -> Option<(TypeInfo, StringLiteralValue)> {
+        let Expr::Attribute(attr_expr) = func else {
+            return None;
+        };
+        if attr_expr.attr.id.as_str() != "get" {
+            return None;
+        }
+        if args.args.len() != 1 {
+            return None;
+        }
+        let Expr::StringLiteral(ExprStringLiteral { value: key, .. }) = &args.args[0] else {
+            return None;
+        };
+        let obj_ty = self.expr_infer_type_info_with_hint(&attr_expr.value, None, errors);
+        if self.is_dict_like(obj_ty.ty()) {
+            Some((obj_ty, key.clone()))
+        } else {
+            None
+        }
+    }
+
+    // Is this type a `TypedDict` or subtype of `dict`, but not `Any`?
+    pub fn is_dict_like(&self, ty: &Type) -> bool {
+        if ty.is_any() {
+            return false;
+        }
+        if ty.is_typed_dict() {
+            return true;
+        }
+        let dict_type = self
+            .stdlib
+            .dict(Type::any_implicit(), Type::any_implicit())
+            .to_type();
+        self.is_subset_eq(ty, &dict_type)
+    }
+
     /// Determine the boolean behavior of a type:
     /// - `Some(true)` or `Some(false)` when it is known to be statically truthy
     ///   or falsey (as determined by some baked in rules for literals

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;
@@ -821,15 +823,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
                         }
@@ -851,15 +855,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
                     }
@@ -884,16 +890,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::GetMethod
+                    && !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);
@@ -939,4 +950,32 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             ),
         }
     }
+
+    /// We only narrow `x.get("key")` if `x` resolves to a `dict`
+    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(&base_ty)
+    }
 }

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,45 @@ impl AtomicNarrowOp {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum FacetOrigin {
+    // This facet is a direct access, like `x.y`, `x[0]`, `x["key"]`
+    Direct,
+    // This facet came from a call to a `get` method, like `x.get("key")`
+    GetMethod,
+}
+
+#[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 +398,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 +419,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()));
@@ -587,6 +613,12 @@ impl NarrowOps {
                     *range,
                 )
             }
+            Some(e @ Expr::Call(call)) if dict_get_subject_for_call_expr(call).is_some() => {
+                // When the guard is something like `x.get("key")`, we narrow it like `x["key"]` if `x` resolves to a dict
+                // in the answers step.
+                // This cannot be a TypeGuard/TypeIs function call, since the first argument is a string literal
+                Self::from_single_narrow_op(e, AtomicNarrowOp::IsTruthy, e.range())
+            }
             Some(Expr::Call(ExprCall {
                 node_index: _,
                 range,
@@ -770,16 +802,63 @@ pub fn identifier_and_chain_prefix_for_expr(expr: &Expr) -> Option<(Identifier,
     f(expr, Vec::new())
 }
 
-fn subject_for_expr(expr: &Expr) -> Option<NarrowingSubject> {
-    identifier_and_chain_for_expr(expr)
-        .map(|(identifier, attr)| NarrowingSubject::Facets(identifier.id, attr))
+// Handle narrowing on `dict.get("key")`. During solving, if the resolved
+// type of the object is not a subtype of `dict`, we will not perform any narrowing.
+fn dict_get_subject_for_call_expr(call_expr: &ExprCall) -> Option<NarrowingSubject> {
+    let func = &call_expr.func;
+    let arguments = &call_expr.arguments;
+    if arguments.keywords.is_empty()
+        && arguments.args.len() == 1
+        && 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) {
+            // x.y.z.get("key")
+            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::GetMethod,
+                },
+            ));
+        } else if let Expr::Name(name) = &*attr.value {
+            // x.get("key")
+            return Some(NarrowingSubject::Facets(
+                name.id.clone(),
+                FacetSubject {
+                    chain: FacetChain::new(Vec1::new(FacetKind::Key(key))),
+                    origin: FacetOrigin::GetMethod,
+                },
+            ));
+        }
+    }
+    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(_) => {
+                if let Some((identifier, facets)) = identifier_and_chain_for_expr(expr) {
+                    res.push(NarrowingSubject::Facets(
+                        identifier.id,
+                        FacetSubject {
+                            chain: facets,
+                            origin: FacetOrigin::Direct,
+                        },
+                    ));
+                }
+            }
+            Expr::Call(call) => {
+                if let Some(subject) = dict_get_subject_for_call_expr(call) {
+                    res.push(subject);
+                }
+            }
             Expr::Named(ExprNamed { target, value, .. }) => {
                 f(target, res);
                 f(value, res);