Keep literal string key types when unpacking a conditional mapping (#5632)

pyrefly/lib/alt/expr.rs

@@ -2412,7 +2412,20 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     }
                 }
                 None => {
-                    let ty = self.expr_infer(&x.value, errors);
+                    // Propagate the key and value hints into the unpacked operand so compound
+                    // expressions keep literal key types; without them the keys widen to `str`.
+                    let ty = if let (Some(key_hint), Some(value_hint)) = (&key_hint, &value_hint) {
+                        let mapping_hint = self.heap.mk_class_type(
+                            self.stdlib.mapping(key_hint.clone(), value_hint.clone()),
+                        );
+                        self.expr_infer_with_hint(
+                            &x.value,
+                            Some(HintRef::soft(&mapping_hint)),
+                            errors,
+                        )
+                    } else {
+                        self.expr_infer(&x.value, errors)
+                    };
                     let mapping_types = iteration_keys
                         .as_ref()
                         .and_then(|_| self.unwrap_mapping_for_unpacking(&ty));

pyrefly/lib/alt/unwrap.rs

@@ -162,7 +162,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
 
     /// Extract key and value types for dictionary unpacking, including structural mappings.
     pub fn unwrap_mapping_for_unpacking(&self, ty: &Type) -> Option<(Type, Type)> {
-        self.unwrap_mapping(ty).or_else(|| {
+        let ordinary = self.unwrap_mapping(ty).or_else(|| {
             let key = self.fresh_var();
             let value = self.fresh_var();
             let Type::ClassDef(mapping_class) = self.try_get_from_export(
@@ -183,7 +183,23 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             } else {
                 None
             }
-        })
+        });
+        if let Some(pair) = ordinary {
+            return Some(pair);
+        }
+        // Some unions of mappings are not recognized by the extraction above. Unpack those as
+        // the union of each member's key and value types, instead of reporting a non-mapping.
+        let Type::Union(union) = ty else {
+            return None;
+        };
+        let mut keys = Vec::with_capacity(union.members.len());
+        let mut values = Vec::with_capacity(union.members.len());
+        for member in &union.members {
+            let (key, value) = self.unwrap_mapping_for_unpacking(member)?;
+            keys.push(key);
+            values.push(value);
+        }
+        Some((self.unions(keys), self.unions(values)))
     }
 
     /// Warning: this returns `Some` if the type is `Any` or a class that extends `Any`

pyrefly/lib/test/dict.rs

@@ -394,3 +394,33 @@ assert_type(frozendict(x=1), frozendict[str, int])
 assert_type(frozendict({"x": 1}, y=2), frozendict[str, int])
 "#,
 );
+
+// https://github.com/pyrefly-org/pyrefly/issues/5632
+testcase!(
+    test_dict_unpack_conditional_literalstring_keys,
+    r#"
+from typing import Literal, LiteralString, Mapping
+
+def direct(value: str) -> Mapping[LiteralString, str]:
+    return {**{"left": value}}
+
+def conditional(value: str, condition: bool) -> Mapping[LiteralString, str]:
+    return {**({"left": value} if condition else {"right": value})}
+
+def without_unpacking(value: str, condition: bool) -> Mapping[LiteralString, str]:
+    return {"left": value} if condition else {"right": value}
+
+def bad_value(value: str, condition: bool) -> Mapping[LiteralString, str]:
+    return {**({"left": value} if condition else {"right": 1})}  # E: Returned type `dict[str, int | str]` is not assignable to declared return type `Mapping[LiteralString, str]`
+
+def union_of_dicts(
+    c: dict[Literal["left"], str] | dict[Literal["right"], str],
+) -> Mapping[LiteralString, str]:
+    return {**c}
+
+def non_literal_key(
+    key: str, value: str, condition: bool
+) -> Mapping[LiteralString, str]:
+    return {**({key: value} if condition else {"right": value})}  # E: Returned type `dict[str, str]` is not assignable to declared return type `Mapping[LiteralString, str]`
+"#,
+);