diff --git a/crates/pyrefly_types/src/stdlib.rs b/crates/pyrefly_types/src/stdlib.rs index a3202794f4..195625889a 100644 --- a/crates/pyrefly_types/src/stdlib.rs +++ b/crates/pyrefly_types/src/stdlib.rs @@ -78,6 +78,7 @@ pub struct Stdlib { dict_keys: StdlibResult<(Class, Arc)>, dict_values: StdlibResult<(Class, Arc)>, mapping: StdlibResult<(Class, Arc)>, + supports_keys_and_get_item: StdlibResult<(Class, Arc)>, set: StdlibResult<(Class, Arc)>, tuple: StdlibResult<(Class, Arc)>, enumerate: StdlibResult<(Class, Arc)>, @@ -281,6 +282,11 @@ impl Stdlib { method_type: lookup_concrete(types, "MethodType"), module_type: lookup_concrete(types, "ModuleType"), mapping: lookup_generic(typing, "Mapping", 2), + supports_keys_and_get_item: lookup_generic( + ModuleName::from_str("_typeshed"), + "SupportsKeysAndGetItem", + 2, + ), enum_meta: lookup_concrete(enum_, "EnumMeta"), protocol_meta: lookup_concrete(typing, "_ProtocolMeta"), enum_flag: lookup_concrete(enum_, "Flag"), @@ -526,6 +532,10 @@ impl Stdlib { &Self::unwrap(&self.mapping).0 } + pub fn supports_keys_and_get_item(&self, key: Type, value: Type) -> ClassType { + Self::apply(&self.supports_keys_and_get_item, vec![key, value]) + } + pub fn set(&self, x: Type) -> ClassType { Self::apply(&self.set, vec![x]) } diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index 4902ef0bbe..2b18a4d6d5 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/pyrefly/lib/alt/expr.rs @@ -1982,7 +1982,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { value_tys.push(field.ty.clone()); } } - } else if let Some((key_t, value_t)) = self.unwrap_mapping(&ty) { + } else if let Some((key_t, value_t)) = self.unwrap_mapping_for_unpacking(&ty) { // Non-anonymous-typed-dict unpacking disables anonymous typed dict creation can_create_anonymous_typed_dict = false; if !key_t.is_error() { diff --git a/pyrefly/lib/alt/unwrap.rs b/pyrefly/lib/alt/unwrap.rs index bde9e52f93..53563e9ead 100644 --- a/pyrefly/lib/alt/unwrap.rs +++ b/pyrefly/lib/alt/unwrap.rs @@ -172,6 +172,23 @@ 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 key = self.fresh_var(); + let value = self.fresh_var(); + let mapping_type = self.heap.mk_class_type( + self.stdlib + .supports_keys_and_get_item(key.to_type(self.heap), value.to_type(self.heap)), + ); + if self.is_subset_eq(ty, &mapping_type) { + Some((self.resolve_var(ty, key), self.resolve_var(ty, value))) + } else { + None + } + }) + } + /// Warning: this returns `Some` if the type is `Any` or a class that extends `Any` pub fn unwrap_awaitable(&self, ty: &Type) -> Option { let var = self.fresh_var(); diff --git a/pyrefly/lib/test/simple.rs b/pyrefly/lib/test/simple.rs index 3af3c9fbf5..9e675a2c96 100644 --- a/pyrefly/lib/test/simple.rs +++ b/pyrefly/lib/test/simple.rs @@ -967,6 +967,65 @@ def test(m: Mapping[str, int]) -> None: "#, ); +testcase!( + test_dict_unpack_duck_mapping, + r#" +from collections.abc import Iterable +from typing import Any, assert_type + +class DuckMapping: + def keys(self) -> Iterable[str]: + return ["a", "b"] + + def __getitem__(self, key: str) -> Any: + return {"a": 1, "b": 2}[key] + +assert_type({**DuckMapping()}, dict[str, Any]) +assert_type({"x": 1, **DuckMapping()}, dict[str, int | Any]) + +class GenericMapping[K, V]: + def keys(self) -> Iterable[K]: ... + def __getitem__(self, key: K) -> V: ... + +def test(mapping: GenericMapping[str, int], union: GenericMapping[str, int] | dict[str, str]): + assert_type({**mapping}, dict[str, int]) + assert_type({"x": 1, **mapping}, dict[str, int]) + assert_type({**union}, dict[str, int | str]) + widened: dict[str, int | str] = {**mapping} + wrong_key: dict[int, int] = {**mapping} # E: `dict[str, int]` is not assignable to `dict[int, int]` + wrong_value: dict[str, str] = {**mapping} # E: `dict[str, int]` is not assignable to `dict[str, str]` + +def test_non_string_keys(mapping: GenericMapping[int, str]): + assert_type({**mapping}, dict[int, str]) +"#, +); + +testcase!( + test_dict_unpack_invalid_duck_mapping, + r#" +from collections.abc import Iterable + +class MissingKeys: + def __getitem__(self, key: str) -> int: ... + +class MissingGetItem: + def keys(self) -> Iterable[str]: ... + +class NonIterableKeys: + def keys(self) -> int: ... + def __getitem__(self, key: str) -> int: ... + +class IncompatibleKey: + def keys(self) -> Iterable[str]: ... + def __getitem__(self, key: int) -> int: ... + +{**MissingKeys()} # E: Expected a mapping, got MissingKeys +{**MissingGetItem()} # E: Expected a mapping, got MissingGetItem +{**NonIterableKeys()} # E: Expected a mapping, got NonIterableKeys +{**IncompatibleKey()} # E: Expected a mapping, got IncompatibleKey +"#, +); + testcase!( test_dict_unpack_subclass, r#"