Fix missing type store for overloads

mypy/checkexpr.py

@@ -2825,6 +2825,7 @@ def infer_overload_return_type(
                 # Return early if possible; otherwise record info, so we can
                 # check for ambiguity due to 'Any' below.
                 if not args_contain_any:
+                    self.chk.store_types(m)
                     return ret_type, infer_type
                 p_infer_type = get_proper_type(infer_type)
                 if isinstance(p_infer_type, CallableType):

test-data/unit/check-generics.test

@@ -1481,6 +1481,30 @@ if int():
     b = f(b)
 [builtins fixtures/list.pyi]
 
+[case testGenericDictWithOverload]
+from typing import Dict, Generic, TypeVar, Any, overload
+T = TypeVar("T")
+
+class Key(Generic[T]): ...
+class CustomDict(dict):
+    @overload  # type: ignore[override]
+    def __setitem__(self, key: Key[T], value: T) -> None: ...
+    @overload
+    def __setitem__(self, key: str, value: Any) -> None: ...
+    def __setitem__(self, key, value):
+        return super().__setitem__(key, value)
+
+def a1(d: Dict[str, Any]) -> None:
+    if (var := d.get("arg")) is None:
+        var = d["arg"] = {}
+        reveal_type(var)  # N: Revealed type is "builtins.dict[Any, Any]"
+
+def a2(d: CustomDict) -> None:
+    if (var := d.get("arg")) is None:
+        var = d["arg"] = {}
+        reveal_type(var)  # N: Revealed type is "builtins.dict[Any, Any]"
+[builtins fixtures/dict.pyi]
+
 
 -- Type variable scoping
 -- ---------------------

test-data/unit/typexport-basic.test

@@ -1236,6 +1236,27 @@ LambdaExpr(10) : def (x: builtins.int) -> builtins.int
 LambdaExpr(12) : def (y: builtins.str) -> builtins.str
 LambdaExpr(13) : def (x: builtins.str) -> builtins.str
 
+[case testExportOverloadArgTypeDict]
+## DictExpr
+from typing import TypeVar, Generic, Any, overload, Dict
+T = TypeVar("T")
+class Key(Generic[T]): ...
+@overload
+def f(x: Key[T], y: T) -> T: ...
+@overload
+def f(x: int, y: Any) -> Any: ...
+def f(x, y): ...
+d: Dict = {}
+d.get(
+  "", {})
+f(
+  2, {})
+[builtins fixtures/dict.pyi]
+[out]
+DictExpr(10) : builtins.dict[Any, Any]
+DictExpr(12) : builtins.dict[Any, Any]
+DictExpr(14) : builtins.dict[Any, Any]
+
 -- TODO
 --
 -- test expressions