Use more precise context for TypedDict plugin errors (#18293)

Fixes #12271

Uses an applicable argument expression as the error context instead of
the overall CallExpr.

**Given:**
```python
# flags: --pretty --show-column-number
from typing import TypedDict

class A(TypedDict):
    x: int

a: A
x.setdefault("y", 123)
x.setdefault("x", "bad")

# Non-TypedDict case for reference
b: dict[str, int]
b.setdefault("x", "bad")
```
**Before:**
```
main.py:8:1: error: TypedDict "A" has no key "y"  [typeddict-item]
    a.setdefault("y", 123)
    ^~~~~~~~~~~~~~~~~~~~~~
main.py:9:1: error: Argument 2 to "setdefault" of "TypedDict" has incompatible type "str"; expected "int"  [typeddict-item]
    a.setdefault("x", "bad")
    ^~~~~~~~~~~~~~~~~~~~~~~~
main.py:13:19: error: Argument 2 to "setdefault" of "MutableMapping" has incompatible type "str"; expected "int"  [arg-type]
    b.setdefault("x", "bad")
                      ^~~~~
Found 3 errors in 1 file (checked 1 source file)
```
**After:**
```
main.py:8:14: error: TypedDict "A" has no key "y"  [typeddict-item]
    a.setdefault("y", 123)
                 ^~~
main.py:9:19: error: Argument 2 to "setdefault" of "TypedDict" has incompatible type "str"; expected "int"  [typeddict-item]
    a.setdefault("x", "bad")
                      ^~~~~
main.py:13:19: error: Argument 2 to "setdefault" of "MutableMapping" has incompatible type "str"; expected "int"  [arg-type]
    b.setdefault("x", "bad")
                      ^~~~~
Found 3 errors in 1 file (checked 1 source file)
```

Author: brianschubert

mypy/plugins/default.py

@@ -304,25 +304,26 @@ def typed_dict_pop_callback(ctx: MethodContext) -> Type:
         and len(ctx.arg_types) >= 1
         and len(ctx.arg_types[0]) == 1
     ):
-        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        key_expr = ctx.args[0][0]
+        keys = try_getting_str_literals(key_expr, ctx.arg_types[0][0])
         if keys is None:
             ctx.api.fail(
                 message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL,
-                ctx.context,
+                key_expr,
                 code=codes.LITERAL_REQ,
             )
             return AnyType(TypeOfAny.from_error)
 
         value_types = []
         for key in keys:
             if key in ctx.type.required_keys:
-                ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, ctx.context)
+                ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, key_expr)
 
             value_type = ctx.type.items.get(key)
             if value_type:
                 value_types.append(value_type)
             else:
-                ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
+                ctx.api.msg.typeddict_key_not_found(ctx.type, key, key_expr)
                 return AnyType(TypeOfAny.from_error)
 
         if len(ctx.args[1]) == 0:
@@ -363,27 +364,29 @@ def typed_dict_setdefault_callback(ctx: MethodContext) -> Type:
         and len(ctx.arg_types[0]) == 1
         and len(ctx.arg_types[1]) == 1
     ):
-        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        key_expr = ctx.args[0][0]
+        keys = try_getting_str_literals(key_expr, ctx.arg_types[0][0])
         if keys is None:
             ctx.api.fail(
                 message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL,
-                ctx.context,
+                key_expr,
                 code=codes.LITERAL_REQ,
             )
             return AnyType(TypeOfAny.from_error)
 
         assigned_readonly_keys = ctx.type.readonly_keys & set(keys)
         if assigned_readonly_keys:
-            ctx.api.msg.readonly_keys_mutated(assigned_readonly_keys, context=ctx.context)
+            ctx.api.msg.readonly_keys_mutated(assigned_readonly_keys, context=key_expr)
 
         default_type = ctx.arg_types[1][0]
+        default_expr = ctx.args[1][0]
 
         value_types = []
         for key in keys:
             value_type = ctx.type.items.get(key)
 
             if value_type is None:
-                ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
+                ctx.api.msg.typeddict_key_not_found(ctx.type, key, key_expr)
                 return AnyType(TypeOfAny.from_error)
 
             # The signature_callback above can't always infer the right signature
@@ -392,7 +395,7 @@ def typed_dict_setdefault_callback(ctx: MethodContext) -> Type:
             # default can be assigned to all key-value pairs we're updating.
             if not is_subtype(default_type, value_type):
                 ctx.api.msg.typeddict_setdefault_arguments_inconsistent(
-                    default_type, value_type, ctx.context
+                    default_type, value_type, default_expr
                 )
                 return AnyType(TypeOfAny.from_error)
 
@@ -409,20 +412,21 @@ def typed_dict_delitem_callback(ctx: MethodContext) -> Type:
         and len(ctx.arg_types) == 1
         and len(ctx.arg_types[0]) == 1
     ):
-        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        key_expr = ctx.args[0][0]
+        keys = try_getting_str_literals(key_expr, ctx.arg_types[0][0])
         if keys is None:
             ctx.api.fail(
                 message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL,
-                ctx.context,
+                key_expr,
                 code=codes.LITERAL_REQ,
             )
             return AnyType(TypeOfAny.from_error)
 
         for key in keys:
             if key in ctx.type.required_keys or key in ctx.type.readonly_keys:
-                ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, ctx.context)
+                ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, key_expr)
             elif key not in ctx.type.items:
-                ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
+                ctx.api.msg.typeddict_key_not_found(ctx.type, key, key_expr)
     return ctx.default_return_type
 
 

test-data/unit/check-columns.test

@@ -227,9 +227,19 @@ class D(TypedDict):
     x: int
 t: D = {'x':
     'y'} # E:5: Incompatible types (expression has type "str", TypedDict item "x" has type "int")
+s: str
 
 if int():
-    del t['y'] # E:5: TypedDict "D" has no key "y"
+    del t[s]    # E:11: Expected TypedDict key to be string literal
+    del t["x"]  # E:11: Key "x" of TypedDict "D" cannot be deleted
+    del t["y"]  # E:11: TypedDict "D" has no key "y"
+
+t.pop(s)    # E:7: Expected TypedDict key to be string literal
+t.pop("y")  # E:7: TypedDict "D" has no key "y"
+
+t.setdefault(s, 123)    # E:14: Expected TypedDict key to be string literal
+t.setdefault("x", "a")  # E:19: Argument 2 to "setdefault" of "TypedDict" has incompatible type "str"; expected "int"
+t.setdefault("y", 123)  # E:14: TypedDict "D" has no key "y"
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 

test-data/unit/check-literal.test

@@ -1909,8 +1909,9 @@ reveal_type(d.get(a_key, u))  # N: Revealed type is "Union[builtins.int, __main_
 reveal_type(d.get(b_key, u))  # N: Revealed type is "Union[builtins.str, __main__.Unrelated]"
 reveal_type(d.get(c_key, u))  # N: Revealed type is "builtins.object"
 
-reveal_type(d.pop(a_key))     # E: Key "a" of TypedDict "Outer" cannot be deleted \
-                              # N: Revealed type is "builtins.int"
+reveal_type(d.pop(a_key))     # N: Revealed type is "builtins.int" \
+                              # E: Key "a" of TypedDict "Outer" cannot be deleted
+
 reveal_type(d.pop(b_key))     # N: Revealed type is "builtins.str"
 d.pop(c_key)                  # E: TypedDict "Outer" has no key "c"
 

test-data/unit/check-typeddict.test

@@ -1747,8 +1747,9 @@ td: Union[TDA, TDB]
 
 reveal_type(td.pop('a'))  # N: Revealed type is "builtins.int"
 reveal_type(td.pop('b'))  # N: Revealed type is "Union[builtins.str, builtins.int]"
-reveal_type(td.pop('c'))  # E: TypedDict "TDA" has no key "c" \
-                          # N: Revealed type is "Union[Any, builtins.int]"
+reveal_type(td.pop('c'))  # N: Revealed type is "Union[Any, builtins.int]" \
+                          # E: TypedDict "TDA" has no key "c"
+
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -2614,8 +2615,9 @@ def func(foo: Union[Foo1, Foo2]):
 
     del foo["missing"]  # E: TypedDict "Foo1" has no key "missing" \
                         # E: TypedDict "Foo2" has no key "missing"
-    del foo[1]  # E: Expected TypedDict key to be string literal \
-                # E: Argument 1 to "__delitem__" has incompatible type "int"; expected "str"
+    del foo[1]  # E: Argument 1 to "__delitem__" has incompatible type "int"; expected "str" \
+                # E: Expected TypedDict key to be string literal
+
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -3726,8 +3728,9 @@ class TP(TypedDict):
     mutable: bool
 
 x: TP
-reveal_type(x.pop("key"))  # E: Key "key" of TypedDict "TP" cannot be deleted \
-                           # N: Revealed type is "builtins.str"
+reveal_type(x.pop("key"))  # N: Revealed type is "builtins.str" \
+                           # E: Key "key" of TypedDict "TP" cannot be deleted
+
 
 x.update({"key": "abc", "other": 1, "mutable": True})  # E: ReadOnly TypedDict keys ("key", "other") TypedDict are mutated
 x.setdefault("key", "abc")  # E: ReadOnly TypedDict key "key" TypedDict is mutated