revert reveal_type changes

Author: randolf-scholz

mypy/checkexpr.py

@@ -1502,7 +1502,7 @@ def check_call_expr_with_callee_type(
     def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type:
         """Type check calling a member expression where the base type is a union."""
         res: list[Type] = []
-        for typ in object_type.relevant_items():
+        for typ in flatten_nested_unions(object_type.relevant_items()):
             # Member access errors are already reported when visiting the member expression.
             with self.msg.filter_errors():
                 item = analyze_member_access(

mypy/checkmember.py

@@ -17,7 +17,6 @@
 from mypy.meet import is_overlapping_types
 from mypy.messages import MessageBuilder
 from mypy.nodes import (
-    ARG_OPT,
     ARG_POS,
     ARG_STAR,
     ARG_STAR2,
@@ -69,7 +68,6 @@
     TypedDictType,
     TypeOfAny,
     TypeType,
-    TypeVarId,
     TypeVarLikeType,
     TypeVarTupleType,
     TypeVarType,
@@ -1406,76 +1404,6 @@ def analyze_typeddict_access(
             fallback=mx.chk.named_type("builtins.function"),
             name=name,
         )
-    elif name == "get":
-        # synthesize TypedDict.get() overloads
-        str_type = mx.chk.named_type("builtins.str")
-        fn_type = mx.chk.named_type("builtins.function")
-        object_type = mx.chk.named_type("builtins.object")
-        type_info = typ.fallback.type
-        # type variable for default value
-        tvar = TypeVarType(
-            "T",
-            f"{type_info.fullname}.get.T",
-            id=TypeVarId(-1, namespace=f"{type_info.fullname}.get"),
-            values=[],
-            upper_bound=object_type,
-            default=AnyType(TypeOfAny.from_omitted_generics),
-        )
-        # generate the overloads
-        overloads: list[CallableType] = []
-        for key, value_type in typ.items.items():
-            key_type = LiteralType(key, fallback=str_type)
-
-            if key in typ.required_keys:
-                # If the key is required, we know it must be present in the TypedDict.
-                # def (K, object=...) -> V
-                overload = CallableType(
-                    arg_types=[key_type, object_type],
-                    arg_kinds=[ARG_POS, ARG_OPT],
-                    arg_names=[None, None],
-                    ret_type=value_type,
-                    fallback=fn_type,
-                    name=name,
-                )
-                overloads.append(overload)
-            else:
-                # The key is not required, but if it is present, we know its type.
-                # def (K) -> V | None   (implicit default)
-                overload = CallableType(
-                    arg_types=[key_type],
-                    arg_kinds=[ARG_POS],
-                    arg_names=[None],
-                    ret_type=UnionType.make_union([value_type, NoneType()]),
-                    fallback=fn_type,
-                    name=name,
-                )
-                overloads.append(overload)
-
-                # def [T](K, T) -> V | T   (explicit default)
-                overload = CallableType(
-                    variables=[tvar],
-                    arg_types=[key_type, tvar],
-                    arg_kinds=[ARG_POS, ARG_POS],
-                    arg_names=[None, None],
-                    ret_type=UnionType.make_union([value_type, tvar]),
-                    fallback=fn_type,
-                    name=name,
-                )
-                overloads.append(overload)
-
-        # finally, add fallback overload when a key is used that is not in the TypedDict
-        # TODO: add support for extra items (PEP 728)
-        # def (str, object=...) -> object
-        fallback_overload = CallableType(
-            arg_types=[str_type, object_type],
-            arg_kinds=[ARG_POS, ARG_OPT],
-            arg_names=[None, None],
-            ret_type=object_type,
-            fallback=fn_type,
-            name=name,
-        )
-        overloads.append(fallback_overload)
-        return Overloaded(overloads)
     return _analyze_member_access(name, typ.fallback, mx, override_info)
 
 

mypy/plugins/default.py

@@ -120,9 +120,9 @@ def get_function_signature_hook(
     def get_method_signature_hook(
         self, fullname: str
     ) -> Callable[[MethodSigContext], FunctionLike] | None:
-        # NOTE: signatures for `__setitem__`, `__delitem__` and `get` are
-        #  defined in checkmember.py/analyze_typeddict_access
-        if fullname in TD_SETDEFAULT_NAMES:
+        if fullname == "typing.Mapping.get":
+            return typed_dict_get_signature_callback
+        elif fullname in TD_SETDEFAULT_NAMES:
             return typed_dict_setdefault_signature_callback
         elif fullname in TD_POP_NAMES:
             return typed_dict_pop_signature_callback
@@ -212,6 +212,46 @@ def get_class_decorator_hook_2(
         return None
 
 
+def typed_dict_get_signature_callback(ctx: MethodSigContext) -> CallableType:
+    """Try to infer a better signature type for TypedDict.get.
+
+    This is used to get better type context for the second argument that
+    depends on a TypedDict value type.
+    """
+    signature = ctx.default_signature
+    if (
+        isinstance(ctx.type, TypedDictType)
+        and len(ctx.args) == 2
+        and len(ctx.args[0]) == 1
+        and isinstance(ctx.args[0][0], StrExpr)
+        and len(signature.arg_types) == 2
+        and len(signature.variables) == 1
+        and len(ctx.args[1]) == 1
+    ):
+        key = ctx.args[0][0].value
+        value_type = get_proper_type(ctx.type.items.get(key))
+        ret_type = signature.ret_type
+        if value_type:
+            default_arg = ctx.args[1][0]
+            if (
+                isinstance(value_type, TypedDictType)
+                and isinstance(default_arg, DictExpr)
+                and len(default_arg.items) == 0
+            ):
+                # Caller has empty dict {} as default for typed dict.
+                value_type = value_type.copy_modified(required_keys=set())
+            # Tweak the signature to include the value type as context. It's
+            # only needed for type inference since there's a union with a type
+            # variable that accepts everything.
+            tv = signature.variables[0]
+            assert isinstance(tv, TypeVarType)
+            return signature.copy_modified(
+                arg_types=[signature.arg_types[0], make_simplified_union([value_type, tv])],
+                ret_type=ret_type,
+            )
+    return signature
+
+
 def typed_dict_get_callback(ctx: MethodContext) -> Type:
     """Infer a precise return type for TypedDict.get with literal first argument."""
     if (

test-data/unit/check-typeddict.test

@@ -1007,12 +1007,12 @@ class D(TypedDict):
     b: NotRequired[str]
 
 def test(d: D) -> None:
-    reveal_type(d.get)  # N: Revealed type is \
-        "Overload(\
-        def (Literal['a'], builtins.object =) -> builtins.int, \
-        def (Literal['b']) -> Union[builtins.str, None], \
-        def [T] (Literal['b'], T`-1) -> Union[builtins.str, T`-1], \
-        def (builtins.str, builtins.object =) -> builtins.object)"
+    reveal_type(d.get)  # N: Revealed type is "Overload(def (k: builtins.str) -> builtins.object, def (builtins.str, builtins.object) -> builtins.object, def [V] (builtins.str, V`4) -> builtins.object)"
+
+
+
+
+
 
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
@@ -1155,21 +1155,29 @@ reveal_type(d.get('x', a)) # N: Revealed type is "Union[builtins.list[builtins.i
 from typing import TypedDict
 D = TypedDict('D', {'x': int, 'y': str})
 d: D
-d.get() # E: All overload variants of "get" require at least one argument \
+d.get() # E: All overload variants of "get" of "Mapping" require at least one argument \
         # N: Possible overload variants: \
-        # N:     def get(Literal['x'], object = ..., /) -> int \
-        # N:     def get(Literal['y'], object = ..., /) -> str \
-        # N:     def get(str, object = ..., /) -> object
-d.get('x', 1, 2) # E: No overload variant of "get" matches argument types "str", "int", "int" \
+        # N:     def get(self, k: str) -> object \
+        # N:     def get(self, str, object, /) -> object \
+        # N:     def [V] get(self, str, V, /) -> object
+d.get('x', 1, 2) # E: No overload variant of "get" of "Mapping" matches argument types "str", "int", "int" \
                  # N: Possible overload variants: \
-                 # N:     def get(Literal['x'], object = ..., /) -> int \
-                 # N:     def get(Literal['y'], object = ..., /) -> str \
-                 # N:     def get(str, object = ..., /) -> object
+                 # N:     def get(self, k: str) -> object \
+                 # N:     def get(self, str, object, /) -> object \
+                 # N:     def [V] get(self, str, Union[int, V], /) -> object
 x = d.get('z')
 reveal_type(x) # N: Revealed type is "builtins.object"
 s = ''
 y = d.get(s)
 reveal_type(y) # N: Revealed type is "builtins.object"
+
+
+
+
+
+
+
+
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 

test-data/unit/fixtures/typing-typeddict.pyi

@@ -56,7 +56,9 @@ class Mapping(Iterable[T], Generic[T, T_co], metaclass=ABCMeta):
     @overload
     def get(self, k: T) -> Optional[T_co]: pass
     @overload
-    def get(self, k: T, default: Union[T_co, V]) -> Union[T_co, V]: pass
+    def get(self, k: T, default: T_co, /) -> Optional[T_co]: pass  # type: ignore[misc]
+    @overload
+    def get(self, k: T, default: V, /) -> Union[T_co, V]: pass
     def values(self) -> Iterable[T_co]: pass  # Approximate return type
     def __len__(self) -> int: ...
     def __contains__(self, arg: object) -> int: pass

test-data/unit/pythoneval.test

@@ -1056,22 +1056,20 @@ def test_not_total(d: D_not_total) -> None:
 _testTypedDictGet.py:8: note: Revealed type is "builtins.int"
 _testTypedDictGet.py:9: note: Revealed type is "builtins.str"
 _testTypedDictGet.py:10: note: Revealed type is "builtins.object"
-_testTypedDictGet.py:11: error: All overload variants of "get" require at least one argument
+_testTypedDictGet.py:11: error: All overload variants of "get" of "Mapping" require at least one argument
 _testTypedDictGet.py:11: note: Possible overload variants:
-_testTypedDictGet.py:11: note:     def get(Literal['x'], object = ..., /) -> int
-_testTypedDictGet.py:11: note:     def get(Literal['y'], object = ..., /) -> str
-_testTypedDictGet.py:11: note:     def get(str, object = ..., /) -> object
+_testTypedDictGet.py:11: note:     def get(self, str, /) -> object
+_testTypedDictGet.py:11: note:     def get(self, str, /, default: object) -> object
+_testTypedDictGet.py:11: note:     def [_T] get(self, str, /, default: _T) -> object
 _testTypedDictGet.py:13: note: Revealed type is "builtins.object"
 _testTypedDictGet.py:16: note: Revealed type is "Union[builtins.int, None]"
 _testTypedDictGet.py:17: note: Revealed type is "Union[builtins.str, None]"
 _testTypedDictGet.py:18: note: Revealed type is "builtins.object"
-_testTypedDictGet.py:19: error: All overload variants of "get" require at least one argument
+_testTypedDictGet.py:19: error: All overload variants of "get" of "Mapping" require at least one argument
 _testTypedDictGet.py:19: note: Possible overload variants:
-_testTypedDictGet.py:19: note:     def get(Literal['x'], /) -> Optional[int]
-_testTypedDictGet.py:19: note:     def [T] get(Literal['x'], T, /) -> Union[int, T]
-_testTypedDictGet.py:19: note:     def get(Literal['y'], /) -> Optional[str]
-_testTypedDictGet.py:19: note:     def [T] get(Literal['y'], T, /) -> Union[str, T]
-_testTypedDictGet.py:19: note:     def get(str, object = ..., /) -> object
+_testTypedDictGet.py:19: note:     def get(self, str, /) -> object
+_testTypedDictGet.py:19: note:     def get(self, str, /, default: object) -> object
+_testTypedDictGet.py:19: note:     def [_T] get(self, str, /, default: _T) -> object
 _testTypedDictGet.py:21: note: Revealed type is "builtins.object"
 
 [case testTypedDictMappingMethods]