Fix type inference for Mapping

docs/changelog.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- Fix type inference for subscriptingon `Mapping`
+
 ## Version 0.10.0 (May 10, 2023)
 
 - Infer the signature for built-in static methods, such as `dict.fromkeys` (#619)

pyanalyze/implementation.py

@@ -17,6 +17,7 @@
 from .signature import (
     ANY_SIGNATURE,
     CallContext,
+    Impl,
     ImplReturn,
     ParameterKind,
     Signature,
@@ -530,6 +531,34 @@ def _dict_setitem_impl(ctx: CallContext) -> ImplReturn:
     return _add_pairs_to_dict(ctx.vars["self"], [pair], ctx, varname)
 
 
+def _mapping_getitem_impl(ctx: CallContext) -> ImplReturn:
+    def inner(key: Value) -> Value:
+        self_value = ctx.vars["self"]
+        if isinstance(self_value, AnnotatedValue):
+            self_value = self_value.value
+        if not _check_dict_key_hashability(key, ctx, "k"):
+            return AnyValue(AnySource.error)
+        if isinstance(self_value, TypedValue):
+            key_type = self_value.get_generic_arg_for_type(
+                collections.abc.Mapping, ctx.visitor, 0
+            )
+            can_assign = key_type.can_assign(key, ctx.visitor)
+            if isinstance(can_assign, CanAssignError):
+                ctx.show_error(
+                    f"Mapping does not accept keys of type {key}",
+                    error_code=ErrorCode.incompatible_argument,
+                    detail=str(can_assign),
+                    arg="key",
+                )
+            return self_value.get_generic_arg_for_type(
+                collections.abc.Mapping, ctx.visitor, 1
+            )
+        else:
+            return AnyValue(AnySource.inference)
+
+    return flatten_unions(inner, ctx.vars["k"])
+
+
 def _dict_getitem_impl(ctx: CallContext) -> ImplReturn:
     def inner(key: Value) -> Value:
         self_value = ctx.vars["self"]
@@ -602,6 +631,22 @@ def inner(key: Value) -> Value:
     return flatten_unions(inner, ctx.vars["k"])
 
 
+def _make_mapping_getitem_sig(typ: type, impl: Impl) -> Signature:
+    return Signature.make(
+        [
+            SigParameter(
+                "self",
+                _POS_ONLY,
+                annotation=GenericValue(typ, [TypeVarValue(K), TypeVarValue(V)]),
+            ),
+            SigParameter("k", _POS_ONLY),
+        ],
+        callable=typ.__getitem__,
+        impl=impl,
+        return_annotation=TypeVarValue(V),
+    )
+
+
 def _dict_get_impl(ctx: CallContext) -> ImplReturn:
     default = ctx.vars["default"]
 
@@ -1605,19 +1650,8 @@ def get_default_argspecs() -> Dict[object, Signature]:
             impl=_dict_setitem_impl,
             return_annotation=KnownValue(None),
         ),
-        Signature.make(
-            [
-                SigParameter(
-                    "self",
-                    _POS_ONLY,
-                    annotation=GenericValue(dict, [TypeVarValue(K), TypeVarValue(V)]),
-                ),
-                SigParameter("k", _POS_ONLY),
-            ],
-            callable=dict.__getitem__,
-            impl=_dict_getitem_impl,
-            return_annotation=TypeVarValue(V),
-        ),
+        _make_mapping_getitem_sig(dict, _dict_getitem_impl),
+        _make_mapping_getitem_sig(collections.abc.Mapping, _mapping_getitem_impl),
         Signature.make(
             [
                 SigParameter("self", _POS_ONLY, annotation=TypedValue(dict)),

pyanalyze/test_implementation.py

@@ -1169,6 +1169,30 @@ def capybara(d: dict, t: Type[int]):
             d[int]
             d[t]
 
+    @assert_passes()
+    def test_mapping(self):
+        from typing import Mapping
+        from typing_extensions import assert_type
+
+        def capybara(m: Mapping[str, int], unannotated):
+            assert_type(m["x"], int)
+            assert_type(m[unannotated], int)
+            m[1]  # E: incompatible_argument
+            m[{}]  # E: unhashable_key
+            m["x"] = 3  # E: unsupported_operation
+
+    @assert_passes()
+    def test_mutable_mapping(self):
+        from typing import MutableMapping
+        from typing_extensions import assert_type
+
+        def capybara(m: MutableMapping[str, int], unannotated):
+            assert_type(m["x"], int)
+            assert_type(m[unannotated], int)
+            m[1]  # E: incompatible_argument
+            m[{}]  # E: unhashable_key
+            m["x"] = 3
+
 
 class TestDictSetItem(TestNameCheckVisitorBase):
     @assert_passes()