diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index 6d00cdd9c6..794ddb71af 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/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)); diff --git a/pyrefly/lib/alt/unwrap.rs b/pyrefly/lib/alt/unwrap.rs index 04e64bdb65..d0839080ba 100644 --- a/pyrefly/lib/alt/unwrap.rs +++ b/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` diff --git a/pyrefly/lib/test/dict.rs b/pyrefly/lib/test/dict.rs index 7e9aa64ede..d1b6109a81 100644 --- a/pyrefly/lib/test/dict.rs +++ b/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]` +"#, +);