More precise return types for TypedDict.get

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/plugins/default.py

@@ -5,7 +5,7 @@
 
 import mypy.errorcodes as codes
 from mypy import message_registry
-from mypy.nodes import DictExpr, IntExpr, StrExpr, UnaryExpr
+from mypy.nodes import DictExpr, Expression, IntExpr, StrExpr, UnaryExpr
 from mypy.plugin import (
     AttributeContext,
     ClassDefContext,
@@ -263,30 +263,40 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
         if keys is None:
             return ctx.default_return_type
 
+        default_type: Type
+        default_arg: Expression | None
+        if len(ctx.arg_types) <= 1 or not ctx.arg_types[1]:
+            default_arg = None
+            default_type = NoneType()
+        elif len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
+            default_arg = ctx.args[1][0]
+            default_type = ctx.arg_types[1][0]
+        else:
+            return ctx.default_return_type
+
         output_types: list[Type] = []
         for key in keys:
-            value_type = get_proper_type(ctx.type.items.get(key))
+            value_type: Type | None = ctx.type.items.get(key)
             if value_type is None:
                 return ctx.default_return_type
 
-            if len(ctx.arg_types) == 1:
+            if key in ctx.type.required_keys:
                 output_types.append(value_type)
-            elif len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
-                default_arg = ctx.args[1][0]
+            else:
+                # HACK to deal with get(key, {})
                 if (
                     isinstance(default_arg, DictExpr)
                     and len(default_arg.items) == 0
-                    and isinstance(value_type, TypedDictType)
+                    and isinstance(vt := get_proper_type(value_type), TypedDictType)
                 ):
-                    # Special case '{}' as the default for a typed dict type.
-                    output_types.append(value_type.copy_modified(required_keys=set()))
+                    output_types.append(vt.copy_modified(required_keys=set()))
                 else:
                     output_types.append(value_type)
-                    output_types.append(ctx.arg_types[1][0])
-
-        if len(ctx.arg_types) == 1:
-            output_types.append(NoneType())
+                    output_types.append(default_type)
 
+        # for nicer reveal_type, put default at the end, if it is present
+        if default_type in output_types:
+            output_types = [t for t in output_types if t != default_type] + [default_type]
         return make_simplified_union(output_types)
     return ctx.default_return_type
 

test-data/unit/check-incremental.test

@@ -7311,3 +7311,204 @@ x = 2
 [out]
 [rechecked bar]
 [stale]
+
+
+[case testIncrementalTypedDictGetMethodTotalFalse]
+import impl
+[file lib.py]
+from typing import TypedDict
+class Unrelated: pass
+D = TypedDict('D', {'x': int, 'y': str}, total=False)
+[file impl.py]
+pass
+[file impl.py.2]
+from typing import Literal
+from lib import D, Unrelated
+d: D
+u: Unrelated
+x: Literal['x']
+y: Literal['y']
+z: Literal['z']
+x_or_y: Literal['x', 'y']
+x_or_z: Literal['x', 'z']
+x_or_y_or_z: Literal['x', 'y', 'z']
+
+# test with literal expression
+reveal_type(d.get('x'))
+reveal_type(d.get('y'))
+reveal_type(d.get('z'))
+reveal_type(d.get('x', u))
+reveal_type(d.get('x', 1))
+reveal_type(d.get('y', None))
+
+# test with literal type / union of literal types with implicit default
+reveal_type(d.get(x))
+reveal_type(d.get(y))
+reveal_type(d.get(z))
+reveal_type(d.get(x_or_y))
+reveal_type(d.get(x_or_z))
+reveal_type(d.get(x_or_y_or_z))
+
+# test with literal type / union of literal types with explicit default
+reveal_type(d.get(x, u))
+reveal_type(d.get(y, u))
+reveal_type(d.get(z, u))
+reveal_type(d.get(x_or_y, u))
+reveal_type(d.get(x_or_z, u))
+reveal_type(d.get(x_or_y_or_z, u))
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+[out]
+[out2]
+tmp/impl.py:13: note: Revealed type is "Union[builtins.int, None]"
+tmp/impl.py:14: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:15: note: Revealed type is "builtins.object"
+tmp/impl.py:16: note: Revealed type is "Union[builtins.int, lib.Unrelated]"
+tmp/impl.py:17: note: Revealed type is "builtins.int"
+tmp/impl.py:18: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:21: note: Revealed type is "Union[builtins.int, None]"
+tmp/impl.py:22: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:23: note: Revealed type is "builtins.object"
+tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str, None]"
+tmp/impl.py:25: note: Revealed type is "builtins.object"
+tmp/impl.py:26: note: Revealed type is "builtins.object"
+tmp/impl.py:29: note: Revealed type is "Union[builtins.int, lib.Unrelated]"
+tmp/impl.py:30: note: Revealed type is "Union[builtins.str, lib.Unrelated]"
+tmp/impl.py:31: note: Revealed type is "builtins.object"
+tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str, lib.Unrelated]"
+tmp/impl.py:33: note: Revealed type is "builtins.object"
+tmp/impl.py:34: note: Revealed type is "builtins.object"
+
+[case testIncrementalTypedDictGetMethodTotalTrue]
+import impl
+[file lib.py]
+from typing import TypedDict
+class Unrelated: pass
+D = TypedDict('D', {'x': int, 'y': str}, total=True)
+[file impl.py]
+pass
+[file impl.py.2]
+from typing import Literal
+from lib import D, Unrelated
+d: D
+u: Unrelated
+x: Literal['x']
+y: Literal['y']
+z: Literal['z']
+x_or_y: Literal['x', 'y']
+x_or_z: Literal['x', 'z']
+x_or_y_or_z: Literal['x', 'y', 'z']
+
+# test with literal expression
+reveal_type(d.get('x'))
+reveal_type(d.get('y'))
+reveal_type(d.get('z'))
+reveal_type(d.get('x', u))
+reveal_type(d.get('x', 1))
+reveal_type(d.get('y', None))
+
+# test with literal type / union of literal types with implicit default
+reveal_type(d.get(x))
+reveal_type(d.get(y))
+reveal_type(d.get(z))
+reveal_type(d.get(x_or_y))
+reveal_type(d.get(x_or_z))
+reveal_type(d.get(x_or_y_or_z))
+
+# test with literal type / union of literal types with explicit default
+reveal_type(d.get(x, u))
+reveal_type(d.get(y, u))
+reveal_type(d.get(z, u))
+reveal_type(d.get(x_or_y, u))
+reveal_type(d.get(x_or_z, u))
+reveal_type(d.get(x_or_y_or_z, u))
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+[out]
+[out2]
+tmp/impl.py:13: note: Revealed type is "builtins.int"
+tmp/impl.py:14: note: Revealed type is "builtins.str"
+tmp/impl.py:15: note: Revealed type is "builtins.object"
+tmp/impl.py:16: note: Revealed type is "builtins.int"
+tmp/impl.py:17: note: Revealed type is "builtins.int"
+tmp/impl.py:18: note: Revealed type is "builtins.str"
+tmp/impl.py:21: note: Revealed type is "builtins.int"
+tmp/impl.py:22: note: Revealed type is "builtins.str"
+tmp/impl.py:23: note: Revealed type is "builtins.object"
+tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str]"
+tmp/impl.py:25: note: Revealed type is "builtins.object"
+tmp/impl.py:26: note: Revealed type is "builtins.object"
+tmp/impl.py:29: note: Revealed type is "builtins.int"
+tmp/impl.py:30: note: Revealed type is "builtins.str"
+tmp/impl.py:31: note: Revealed type is "builtins.object"
+tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str]"
+tmp/impl.py:33: note: Revealed type is "builtins.object"
+tmp/impl.py:34: note: Revealed type is "builtins.object"
+
+
+[case testIncrementalTypedDictGetMethodTotalMixed]
+import impl
+[file lib.py]
+from typing import TypedDict
+from typing_extensions import Required, NotRequired
+class Unrelated: pass
+D = TypedDict('D', {'x': Required[int], 'y': NotRequired[str]})
+[file impl.py]
+pass
+[file impl.py.2]
+from typing import Literal
+from lib import D, Unrelated
+d: D
+u: Unrelated
+x: Literal['x']
+y: Literal['y']
+z: Literal['z']
+x_or_y: Literal['x', 'y']
+x_or_z: Literal['x', 'z']
+x_or_y_or_z: Literal['x', 'y', 'z']
+
+# test with literal expression
+reveal_type(d.get('x'))
+reveal_type(d.get('y'))
+reveal_type(d.get('z'))
+reveal_type(d.get('x', u))
+reveal_type(d.get('x', 1))
+reveal_type(d.get('y', None))
+
+# test with literal type / union of literal types with implicit default
+reveal_type(d.get(x))
+reveal_type(d.get(y))
+reveal_type(d.get(z))
+reveal_type(d.get(x_or_y))
+reveal_type(d.get(x_or_z))
+reveal_type(d.get(x_or_y_or_z))
+
+# test with literal type / union of literal types with explicit default
+reveal_type(d.get(x, u))
+reveal_type(d.get(y, u))
+reveal_type(d.get(z, u))
+reveal_type(d.get(x_or_y, u))
+reveal_type(d.get(x_or_z, u))
+reveal_type(d.get(x_or_y_or_z, u))
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-typeddict.pyi]
+[out]
+[out2]
+tmp/impl.py:13: note: Revealed type is "builtins.int"
+tmp/impl.py:14: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:15: note: Revealed type is "builtins.object"
+tmp/impl.py:16: note: Revealed type is "builtins.int"
+tmp/impl.py:17: note: Revealed type is "builtins.int"
+tmp/impl.py:18: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:21: note: Revealed type is "builtins.int"
+tmp/impl.py:22: note: Revealed type is "Union[builtins.str, None]"
+tmp/impl.py:23: note: Revealed type is "builtins.object"
+tmp/impl.py:24: note: Revealed type is "Union[builtins.int, builtins.str, None]"
+tmp/impl.py:25: note: Revealed type is "builtins.object"
+tmp/impl.py:26: note: Revealed type is "builtins.object"
+tmp/impl.py:29: note: Revealed type is "builtins.int"
+tmp/impl.py:30: note: Revealed type is "Union[builtins.str, lib.Unrelated]"
+tmp/impl.py:31: note: Revealed type is "builtins.object"
+tmp/impl.py:32: note: Revealed type is "Union[builtins.int, builtins.str, lib.Unrelated]"
+tmp/impl.py:33: note: Revealed type is "builtins.object"
+tmp/impl.py:34: note: Revealed type is "builtins.object"

test-data/unit/check-literal.test

@@ -1884,7 +1884,7 @@ reveal_type(d[a_key])         # N: Revealed type is "builtins.int"
 reveal_type(d[b_key])         # N: Revealed type is "builtins.str"
 d[c_key]                      # E: TypedDict "Outer" has no key "c"
 
-reveal_type(d.get(a_key, u))  # N: Revealed type is "Union[builtins.int, __main__.Unrelated]"
+reveal_type(d.get(a_key, u))  # N: Revealed type is "builtins.int"
 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"
 
@@ -1928,7 +1928,7 @@ u: Unrelated
 reveal_type(a[int_key_good])         # N: Revealed type is "builtins.int"
 reveal_type(b[int_key_good])         # N: Revealed type is "builtins.int"
 reveal_type(c[str_key_good])         # N: Revealed type is "builtins.int"
-reveal_type(c.get(str_key_good, u))  # N: Revealed type is "Union[builtins.int, __main__.Unrelated]"
+reveal_type(c.get(str_key_good, u))  # N: Revealed type is "builtins.int"
 reveal_type(c.get(str_key_bad, u))   # N: Revealed type is "builtins.object"
 
 a[int_key_bad]                       # E: Tuple index out of range
@@ -1993,8 +1993,8 @@ optional_keys: Literal["d", "e"]
 bad_keys: Literal["a", "bad"]
 
 reveal_type(test[good_keys])                      # N: Revealed type is "Union[__main__.A, __main__.B]"
-reveal_type(test.get(good_keys))                  # N: Revealed type is "Union[__main__.A, __main__.B, None]"
-reveal_type(test.get(good_keys, 3))               # N: Revealed type is "Union[__main__.A, Literal[3]?, __main__.B]"
+reveal_type(test.get(good_keys))                  # N: Revealed type is "Union[__main__.A, __main__.B]"
+reveal_type(test.get(good_keys, 3))               # N: Revealed type is "Union[__main__.A, __main__.B]"
 reveal_type(test.pop(optional_keys))              # N: Revealed type is "Union[__main__.D, __main__.E]"
 reveal_type(test.pop(optional_keys, 3))           # N: Revealed type is "Union[__main__.D, __main__.E, Literal[3]?]"
 reveal_type(test.setdefault(good_keys, AAndB()))  # N: Revealed type is "Union[__main__.A, __main__.B]"
@@ -2037,15 +2037,18 @@ class D2(TypedDict):
     d: D
 
 x: Union[D1, D2]
-bad_keys: Literal['a', 'b', 'c', 'd']
 good_keys: Literal['b', 'c']
+mixed_keys: Literal['a', 'b', 'c', 'd']
+bad_keys: Literal['e', 'f']
 
-x[bad_keys]         # E: TypedDict "D1" has no key "d" \
+x[mixed_keys]         # E: TypedDict "D1" has no key "d" \
                     # E: TypedDict "D2" has no key "a"
 
 reveal_type(x[good_keys])           # N: Revealed type is "Union[__main__.B, __main__.C]"
-reveal_type(x.get(good_keys))       # N: Revealed type is "Union[__main__.B, __main__.C, None]"
-reveal_type(x.get(good_keys, 3))    # N: Revealed type is "Union[__main__.B, Literal[3]?, __main__.C]"
+reveal_type(x.get(good_keys))       # N: Revealed type is "Union[__main__.B, __main__.C]"
+reveal_type(x.get(good_keys, 3))    # N: Revealed type is "Union[__main__.B, __main__.C]"
+reveal_type(x.get(mixed_keys))      # N: Revealed type is "builtins.object"
+reveal_type(x.get(mixed_keys, 3))   # N: Revealed type is "builtins.object"
 reveal_type(x.get(bad_keys))        # N: Revealed type is "builtins.object"
 reveal_type(x.get(bad_keys, 3))     # N: Revealed type is "builtins.object"
 

test-data/unit/check-recursive-types.test

@@ -690,10 +690,11 @@ class TD(TypedDict, total=False):
     y: TD
 
 td: TD
+reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...}), None]"
 td["y"] = {"x": 0, "y": {}}
 td["y"] = {"x": 0, "y": {"x": 0, "y": 42}}  # E: Incompatible types (expression has type "int", TypedDict item "y" has type "TD")
 
-reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...})}), None]"
+reveal_type(td.get("y"))  # N: Revealed type is "Union[TypedDict('__main__.TD', {'x'?: builtins.int, 'y'?: ...}), None]"
 s: str = td.get("y")  # E: Incompatible types in assignment (expression has type "Optional[TD]", variable has type "str")
 
 td.update({"x": 0, "y": {"x": 1, "y": {}}})