Support mappings that do not inherit from collections.abc.Mapping (#…

…501)
quora · Mar 12, 2022 · cfa1d63 · cfa1d63
1 parent 6d8b960
commit cfa1d63
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 3 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Support mappings that do not inherit from `collections.abc.Mapping`
+  (#501)
 - Improve type inference for calls to `set()`, `list()`, and
   `tuple()` with union arguments (#500)
 - Remove special-cased signatured for `sorted()` (#498)

diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
@@ -1518,6 +1518,25 @@ def capybara(
                 ),
             )
 
+    @assert_passes()
+    def test_minimal_mapping(self):
+        from typing import List
+
+        class MyMapping:
+            def keys(self) -> List[bool]:
+                raise NotImplementedError
+
+            def __getitem__(self, key: bool) -> float:
+                raise NotImplementedError
+
+        def capybara(m: MyMapping):
+            assert_is_value(
+                {**m},
+                DictIncompleteValue(
+                    dict, [KVPair(TypedValue(bool), TypedValue(float), is_many=True)]
+                ),
+            )
+
     @assert_passes()
     def test_iterable_unpacking(self):
         def capybara(x):

diff --git a/pyanalyze/value.py b/pyanalyze/value.py
@@ -2262,11 +2262,25 @@ def concrete_values_from_iterable(
 
 K = TypeVar("K")
 V = TypeVar("V")
-MappingValue = GenericValue(collections.abc.Mapping, [TypeVarValue(K), TypeVarValue(V)])
 
 EMPTY_DICTS = (KnownValue({}), DictIncompleteValue(dict, []))
 
 
+# This is all the runtime requires in places like {**k}
+class CustomMapping(Protocol[K, V]):
+    def keys(self) -> Iterable[K]:
+        raise NotImplementedError
+
+    def __getitem__(self, __key: K) -> V:
+        raise NotImplementedError
+
+
+NominalMappingValue = GenericValue(
+    collections.abc.Mapping, [TypeVarValue(K), TypeVarValue(V)]
+)
+ProtocolMappingValue = GenericValue(CustomMapping, [TypeVarValue(K), TypeVarValue(V)])
+
+
 def kv_pairs_from_mapping(
     value_val: Value, ctx: CanAssignContext
 ) -> Union[Sequence[KVPair], CanAssignError]:
@@ -2296,9 +2310,14 @@ def kv_pairs_from_mapping(
             for key, (required, value) in value_val.items.items()
         ]
     else:
-        can_assign = get_tv_map(MappingValue, value_val, ctx)
+        # Ideally we should only need to check ProtocolMappingValue, but if
+        # we do that we can't infer the right types for dict, so try the
+        # nominal Mapping first.
+        can_assign = get_tv_map(NominalMappingValue, value_val, ctx)
         if isinstance(can_assign, CanAssignError):
-            return can_assign
+            can_assign = get_tv_map(ProtocolMappingValue, value_val, ctx)
+            if isinstance(can_assign, CanAssignError):
+                return can_assign
         key_type = can_assign.get(K, AnyValue(AnySource.generic_argument))
         value_type = can_assign.get(V, AnyValue(AnySource.generic_argument))
         return [KVPair(key_type, value_type, is_many=True)]