From 21a97054265909d8c943d8ba9d7e572c3567cb40 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Sat, 20 Jan 2024 17:50:49 +0100
Subject: [PATCH] Fix missing type store for overloads

---
 mypy/checkexpr.py                   |  1 +
 test-data/unit/check-generics.test  | 24 ++++++++++++++++++++++++
 test-data/unit/typexport-basic.test | 21 +++++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index fbedd95e8fd2..be9ded12fb4b 100644
--- a/mypy/checkexpr.py
+++ b/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):
diff --git a/test-data/unit/check-generics.test b/test-data/unit/check-generics.test
index e2f65ed39c1e..116bb12efc60 100644
--- a/test-data/unit/check-generics.test
+++ b/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
 -- ---------------------
diff --git a/test-data/unit/typexport-basic.test b/test-data/unit/typexport-basic.test
index c4c3a1d36f83..d78cf0f179f2 100644
--- a/test-data/unit/typexport-basic.test
+++ b/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