Repository: python/mypy
Pull request: Add union math for intelligent indexing
This pull request hacks on supports for union math when using Unions of Literals to index into tuples, NamedTuples, and TypedDicts.

It also fixes a bug I apparently introduced. Currently, mypy correctly reports an error with this code:

    class Test(TypedDict):
        foo: int

    t: Test

    # Error: int can't be assigned a str value
    t.setdefault("foo", "unrelated value")

...but does not report an error with:

    key: Literal["foo"]
    t.setdefault(key, "unrelated value")

This diff should make mypy report an error in both cases.

Resolves https://github.com/python/mypy/issues/6262.

Problem statement:
Using `Literal` typed variables as keys for TypedDict causes "TypedDict key must be a string literal"
Please provide more information to help us understand the issue:

* Are you reporting a bug, or opening a feature request?
Bug

* Please insert below the code you are checking with mypy,
  or a mock-up repro if the source is private. We would appreciate
  if you try to simplify your case to a minimal repro.

```
from mypy_extensions import TypedDict
from typing_extensions import Literal

class Test(TypedDict):
  a: str
  b: str

TestKey = Literal['a', 'b']

def read(test: Test, key: TestKey):
  return test[key]
```

* What is the actual behavior/output?
mypy.py:13: error: TypedDict key must be a string literal; expected one of ('a', 'b')

* What is the behavior/output you expect?
No type check errors

* What are the versions of mypy and Python you are using?
mypy -- 0.660
python -- 3.6.3

  Do you see the same issue after installing mypy from Git master?
Yes

* What are the mypy flags you are using? (For example --strict-optional)
None

* If mypy crashed with a traceback, please paste
  the full traceback below.

(You can freely edit this text, please remove all the lines
you believe are unnecessary.)

Hints from issue discussion:
Thanks for the bug report!

What's happening in this case is that mypy doesn't yet understand how to handle indexing TypedDicts using Unions of Literals -- it can currently handle indexing using only Literals. It probably should, though. I'll look into making a fix when I find some time later this week/next week.
Thanks!  This is very useful for typing JSON blobs, and having functions that extract information from the JSON blobs using specific keys.

It would also be nice to extract a keys type from a `TypedDict`, like so :-)

```
class Test(TypedDict):
  a: str
  b: str

TestKey = Test.keys()  # same as Literal['a', 'b']
```

Commit message: Add test where we also try a union of typed dicts

Patch to review:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index eab335a47933..8ef01cdfcc8e 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -2,6 +2,7 @@
 
 from collections import OrderedDict
 from contextlib import contextmanager
+import itertools
 from typing import (
     cast, Dict, Set, List, Tuple, Callable, Union, Optional, Sequence, Iterator
 )
@@ -2547,15 +2548,18 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr,
             if isinstance(index, SliceExpr):
                 return self.visit_tuple_slice_helper(left_type, index)
 
-            n = self._get_value(index)
-            if n is not None:
-                if n < 0:
-                    n += len(left_type.items)
-                if 0 <= n < len(left_type.items):
-                    return left_type.items[n]
-                else:
-                    self.chk.fail(message_registry.TUPLE_INDEX_OUT_OF_RANGE, e)
-                    return AnyType(TypeOfAny.from_error)
+            ns = self._get_values(index)
+            if ns is not None:
+                out = []
+                for n in ns:
+                    if n < 0:
+                        n += len(left_type.items)
+                    if 0 <= n < len(left_type.items):
+                        out.append(left_type.items[n])
+                    else:
+                        self.chk.fail(message_registry.TUPLE_INDEX_OUT_OF_RANGE, e)
+                        return AnyType(TypeOfAny.from_error)
+                return UnionType.make_simplified_union(out)
             else:
                 return self.nonliteral_tuple_index_helper(left_type, index)
         elif isinstance(left_type, TypedDictType):
@@ -2571,26 +2575,32 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr,
             return result
 
     def visit_tuple_slice_helper(self, left_type: TupleType, slic: SliceExpr) -> Type:
-        begin = None
-        end = None
-        stride = None
+        begin = [None]   # type: Sequence[Optional[int]]
+        end = [None]     # type: Sequence[Optional[int]]
+        stride = [None]  # type: Sequence[Optional[int]]
 
         if slic.begin_index:
-            begin = self._get_value(slic.begin_index)
-            if begin is None:
+            begin_raw = self._get_values(slic.begin_index)
+            if begin_raw is None:
                 return self.nonliteral_tuple_index_helper(left_type, slic)
+            begin = begin_raw
 
         if slic.end_index:
-            end = self._get_value(slic.end_index)
-            if end is None:
+            end_raw = self._get_values(slic.end_index)
+            if end_raw is None:
                 return self.nonliteral_tuple_index_helper(left_type, slic)
+            end = end_raw
 
         if slic.stride:
-            stride = self._get_value(slic.stride)
-            if stride is None:
+            stride_raw = self._get_values(slic.stride)
+            if stride_raw is None:
                 return self.nonliteral_tuple_index_helper(left_type, slic)
+            stride = stride_raw
 
-        return left_type.slice(begin, stride, end)
+        items = []  # type: List[Type]
+        for b, e, s in itertools.product(begin, end, stride):
+            items.append(left_type.slice(b, e, s))
+        return UnionType.make_simplified_union(items)
 
     def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression) -> Type:
         index_type = self.accept(index)
@@ -2607,40 +2617,59 @@ def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression)
             else:
                 return union
 
-    def _get_value(self, index: Expression) -> Optional[int]:
+    def _get_values(self, index: Expression) -> Optional[List[int]]:
         if isinstance(index, IntExpr):
-            return index.value
+            return [index.value]
         elif isinstance(index, UnaryExpr):
             if index.op == '-':
                 operand = index.expr
                 if isinstance(operand, IntExpr):
-                    return -1 * operand.value
+                    return [-1 * operand.value]
         typ = self.accept(index)
         if isinstance(typ, Instance) and typ.last_known_value is not None:
             typ = typ.last_known_value
         if isinstance(typ, LiteralType) and isinstance(typ.value, int):
-            return typ.value
+            return [typ.value]
+        if isinstance(typ, UnionType):
+            out = []
+            for item in typ.items:
+                if isinstance(item, LiteralType) and isinstance(item.value, int):
+                    out.append(item.value)
+                else:
+                    return None
+            return out
         return None
 
     def visit_typeddict_index_expr(self, td_type: TypedDictType, index: Expression) -> Type:
         if isinstance(index, (StrExpr, UnicodeExpr)):
-            item_name = index.value
+            key_names = [index.value]
         else:
             typ = self.accept(index)
-            if isinstance(typ, Instance) and typ.last_known_value is not None:
-                typ = typ.last_known_value
-
-            if isinstance(typ, LiteralType) and isinstance(typ.value, str):
-                item_name = typ.value
+            if isinstance(typ, UnionType):
+                key_types = typ.items
             else:
-                self.msg.typeddict_key_must_be_string_literal(td_type, index)
-                return AnyType(TypeOfAny.from_error)
+                key_types = [typ]
 
-        item_type = td_type.items.get(item_name)
-        if item_type is None:
-            self.msg.typeddict_key_not_found(td_type, item_name, index)
-            return AnyType(TypeOfAny.from_error)
-        return item_type
+            key_names = []
+            for key_type in key_types:
+                if isinstance(key_type, Instance) and key_type.last_known_value is not None:
+                    key_type = key_type.last_known_value
+
+                if isinstance(key_type, LiteralType) and isinstance(key_type.value, str):
+                    key_names.append(key_type.value)
+                else:
+                    self.msg.typeddict_key_must_be_string_literal(td_type, index)
+                    return AnyType(TypeOfAny.from_error)
+
+        value_types = []
+        for key_name in key_names:
+            value_type = td_type.items.get(key_name)
+            if value_type is None:
+                self.msg.typeddict_key_not_found(td_type, key_name, index)
+                return AnyType(TypeOfAny.from_error)
+            else:
+                value_types.append(value_type)
+        return UnionType.make_simplified_union(value_types)
 
     def visit_enum_index_expr(self, enum_type: TypeInfo, index: Expression,
                               context: Context) -> Type:
diff --git a/mypy/messages.py b/mypy/messages.py
index 9cb279a2d7d0..c9c415dd4ecf 100644
--- a/mypy/messages.py
+++ b/mypy/messages.py
@@ -1170,6 +1170,14 @@ def typeddict_key_cannot_be_deleted(
             self.fail("Key '{}' of TypedDict {} cannot be deleted".format(
                 item_name, self.format(typ)), context)
 
+    def typeddict_setdefault_arguments_inconsistent(
+            self,
+            default: Type,
+            expected: Type,
+            context: Context) -> None:
+        msg = 'Argument 2 to "setdefault" of "TypedDict" has incompatible type {}; expected {}'
+        self.fail(msg.format(self.format(default), self.format(expected)), context)
+
     def type_arguments_not_allowed(self, context: Context) -> None:
         self.fail('Parameterized generics cannot be used with class or instance checks', context)
 
diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py
index adcef71ab5ba..b8bfeb02a623 100644
--- a/mypy/plugins/common.py
+++ b/mypy/plugins/common.py
@@ -6,7 +6,7 @@
 )
 from mypy.plugin import ClassDefContext
 from mypy.semanal import set_callable_name
-from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType, Instance
+from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType, Instance, UnionType
 from mypy.typevars import fill_typevars
 from mypy.util import get_unique_redefinition_name
 
@@ -129,18 +129,34 @@ def add_method(
     info.defn.defs.body.append(func)
 
 
-def try_getting_str_literal(expr: Expression, typ: Type) -> Optional[str]:
-    """If this expression is a string literal, or if the corresponding type
-    is something like 'Literal["some string here"]', returns the underlying
-    string value. Otherwise, returns None."""
+def try_getting_str_literals(expr: Expression, typ: Type) -> Optional[List[str]]:
+    """If the given expression or type corresponds to a string literal
+    or a union of string literals, returns a list of the underlying strings.
+    Otherwise, returns None.
+
+    Specifically, this function is guaranteed to return a list with
+    one or more strings if one one the following is true:
+
+    1. 'expr' is a StrExpr
+    2. 'typ' is a LiteralType containing a string
+    3. 'typ' is a UnionType containing only LiteralType of strings
+    """
+    if isinstance(expr, StrExpr):
+        return [expr.value]
+
     if isinstance(typ, Instance) and typ.last_known_value is not None:
-        typ = typ.last_known_value
-
-    if isinstance(typ, LiteralType) and typ.fallback.type.fullname() == 'builtins.str':
-        val = typ.value
-        assert isinstance(val, str)
-        return val
-    elif isinstance(expr, StrExpr):
-        return expr.value
+        possible_literals = [typ.last_known_value]  # type: List[Type]
+    elif isinstance(typ, UnionType):
+        possible_literals = typ.items
     else:
-        return None
+        possible_literals = [typ]
+
+    strings = []
+    for lit in possible_literals:
+        if isinstance(lit, LiteralType) and lit.fallback.type.fullname() == 'builtins.str':
+            val = lit.value
+            assert isinstance(val, str)
+            strings.append(val)
+        else:
+            return None
+    return strings
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
index fb0177267257..4bfd9129188b 100644
--- a/mypy/plugins/default.py
+++ b/mypy/plugins/default.py
@@ -6,11 +6,12 @@
 from mypy.plugin import (
     Plugin, FunctionContext, MethodContext, MethodSigContext, AttributeContext, ClassDefContext
 )
-from mypy.plugins.common import try_getting_str_literal
+from mypy.plugins.common import try_getting_str_literals
 from mypy.types import (
     Type, Instance, AnyType, TypeOfAny, CallableType, NoneType, UnionType, TypedDictType,
     TypeVarType, TPDICT_FB_NAMES
 )
+from mypy.subtypes import is_subtype
 
 
 class DefaultPlugin(Plugin):
@@ -176,26 +177,34 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
     if (isinstance(ctx.type, TypedDictType)
             and len(ctx.arg_types) >= 1
             and len(ctx.arg_types[0]) == 1):
-        key = try_getting_str_literal(ctx.args[0][0], ctx.arg_types[0][0])
-        if key is None:
+        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        if keys is None:
             return ctx.default_return_type
 
-        value_type = ctx.type.items.get(key)
-        if value_type:
+        output_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)
+                return AnyType(TypeOfAny.from_error)
+
             if len(ctx.arg_types) == 1:
-                return UnionType.make_simplified_union([value_type, NoneType()])
+                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]
                 if (isinstance(default_arg, DictExpr) and len(default_arg.items) == 0
                         and isinstance(value_type, TypedDictType)):
                     # Special case '{}' as the default for a typed dict type.
-                    return value_type.copy_modified(required_keys=set())
+                    output_types.append(value_type.copy_modified(required_keys=set()))
                 else:
-                    return UnionType.make_simplified_union([value_type, ctx.arg_types[1][0]])
-        else:
-            ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
-            return AnyType(TypeOfAny.from_error)
+                    output_types.append(value_type)
+                    output_types.append(ctx.arg_types[1][0])
+
+        if len(ctx.arg_types) == 1:
+            output_types.append(NoneType())
+
+        return UnionType.make_simplified_union(output_types)
     return ctx.default_return_type
 
 
@@ -233,23 +242,28 @@ def typed_dict_pop_callback(ctx: MethodContext) -> Type:
     if (isinstance(ctx.type, TypedDictType)
             and len(ctx.arg_types) >= 1
             and len(ctx.arg_types[0]) == 1):
-        key = try_getting_str_literal(ctx.args[0][0], ctx.arg_types[0][0])
-        if key is None:
+        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        if keys is None:
             ctx.api.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL, ctx.context)
             return AnyType(TypeOfAny.from_error)
 
-        if key in ctx.type.required_keys:
-            ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, ctx.context)
-        value_type = ctx.type.items.get(key)
-        if value_type:
-            if len(ctx.args[1]) == 0:
-                return value_type
-            elif (len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1
-                  and len(ctx.args[1]) == 1):
-                return UnionType.make_simplified_union([value_type, ctx.arg_types[1][0]])
-        else:
-            ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
-            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)
+
+            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)
+                return AnyType(TypeOfAny.from_error)
+
+        if len(ctx.args[1]) == 0:
+            return UnionType.make_simplified_union(value_types)
+        elif (len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1
+              and len(ctx.args[1]) == 1):
+            return UnionType.make_simplified_union([*value_types, ctx.arg_types[1][0]])
     return ctx.default_return_type
 
 
@@ -278,18 +292,35 @@ def typed_dict_setdefault_callback(ctx: MethodContext) -> Type:
     """Type check TypedDict.setdefault and infer a precise return type."""
     if (isinstance(ctx.type, TypedDictType)
             and len(ctx.arg_types) == 2
-            and len(ctx.arg_types[0]) == 1):
-        key = try_getting_str_literal(ctx.args[0][0], ctx.arg_types[0][0])
-        if key is None:
+            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])
+        if keys is None:
             ctx.api.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL, ctx.context)
             return AnyType(TypeOfAny.from_error)
 
-        value_type = ctx.type.items.get(key)
-        if value_type:
-            return value_type
-        else:
-            ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
-            return AnyType(TypeOfAny.from_error)
+        default_type = ctx.arg_types[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)
+                return AnyType(TypeOfAny.from_error)
+
+            # The signature_callback above can't always infer the right signature
+            # (e.g. when the expression is a variable that happens to be a Literal str)
+            # so we need to handle the check ourselves here and make sure the provided
+            # 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)
+                return AnyType(TypeOfAny.from_error)
+
+            value_types.append(value_type)
+
+        return UnionType.make_simplified_union(value_types)
     return ctx.default_return_type
 
 
@@ -304,15 +335,16 @@ def typed_dict_delitem_callback(ctx: MethodContext) -> Type:
     if (isinstance(ctx.type, TypedDictType)
             and len(ctx.arg_types) == 1
             and len(ctx.arg_types[0]) == 1):
-        key = try_getting_str_literal(ctx.args[0][0], ctx.arg_types[0][0])
-        if key is None:
+        keys = try_getting_str_literals(ctx.args[0][0], ctx.arg_types[0][0])
+        if keys is None:
             ctx.api.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL, ctx.context)
             return AnyType(TypeOfAny.from_error)
 
-        if key in ctx.type.required_keys:
-            ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, ctx.context)
-        elif key not in ctx.type.items:
-            ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
+        for key in keys:
+            if key in ctx.type.required_keys:
+                ctx.api.msg.typeddict_key_cannot_be_deleted(ctx.type, key, ctx.context)
+            elif key not in ctx.type.items:
+                ctx.api.msg.typeddict_key_not_found(ctx.type, key, ctx.context)
     return ctx.default_return_type
 
 
diff --git a/mypy/types.py b/mypy/types.py
index 4b688386bcd1..f46844471703 100644
--- a/mypy/types.py
+++ b/mypy/types.py
@@ -1278,8 +1278,8 @@ def copy_modified(self, *, fallback: Optional[Instance] = None,
             items = self.items
         return TupleType(items, fallback, self.line, self.column)
 
-    def slice(self, begin: Optional[int], stride: Optional[int],
-              end: Optional[int]) -> 'TupleType':
+    def slice(self, begin: Optional[int], end: Optional[int],
+              stride: Optional[int]) -> 'TupleType':
         return TupleType(self.items[begin:end:stride], self.partial_fallback,
                          self.line, self.column, self.implicit)
 
diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test
index b2cffb02f1f6..26ae4ff215ed 100644
--- a/test-data/unit/check-literal.test
+++ b/test-data/unit/check-literal.test
@@ -2221,6 +2221,88 @@ c.get(str_key_bad, u)                # E: TypedDict "MyDict" has no key 'missing
 [typing fixtures/typing-full.pyi]
 [out]
 
+[case testLiteralIntelligentIndexingTupleUnions]
+from typing import Tuple, NamedTuple
+from typing_extensions import Literal
+
+class A: pass
+class B: pass
+class C: pass
+class D: pass
+class E: pass
+
+idx1: Literal[1, 2]
+idx2: Literal[3, 4]
+idx_bad: Literal[1, 20]
+
+tup1: Tuple[A, B, C, D, E]
+Tup2Class = NamedTuple('Tup2Class', [('a', A), ('b', B), ('c', C), ('d', D), ('e', E)])
+tup2: Tup2Class
+
+reveal_type(tup1[idx1])         # N: Revealed type is 'Union[__main__.B, __main__.C]'
+reveal_type(tup1[idx1:idx2])    # N: Revealed type is 'Union[Tuple[__main__.B, __main__.C], Tuple[__main__.B, __main__.C, __main__.D], Tuple[__main__.C], Tuple[__main__.C, __main__.D]]'
+reveal_type(tup1[0::idx1])      # N: Revealed type is 'Union[Tuple[__main__.A, __main__.B, __main__.C, __main__.D, __main__.E], Tuple[__main__.A, __main__.C, __main__.E]]'
+tup1[idx_bad]                   # E: Tuple index out of range
+
+reveal_type(tup2[idx1])         # N: Revealed type is 'Union[__main__.B, __main__.C]'
+reveal_type(tup2[idx1:idx2])    # N: Revealed type is 'Union[Tuple[__main__.B, __main__.C, fallback=__main__.Tup2Class], Tuple[__main__.B, __main__.C, __main__.D, fallback=__main__.Tup2Class], Tuple[__main__.C, fallback=__main__.Tup2Class], Tuple[__main__.C, __main__.D, fallback=__main__.Tup2Class]]'
+reveal_type(tup2[0::idx1])      # N: Revealed type is 'Union[Tuple[__main__.A, __main__.B, __main__.C, __main__.D, __main__.E, fallback=__main__.Tup2Class], Tuple[__main__.A, __main__.C, __main__.E, fallback=__main__.Tup2Class]]'
+tup2[idx_bad]                   # E: Tuple index out of range
+[builtins fixtures/slice.pyi]
+[out]
+
+[case testLiteralIntelligentIndexingTypedDictUnions]
+from typing_extensions import Literal, Final
+from mypy_extensions import TypedDict
+
+class A: pass
+class B: pass
+class C: pass
+class D: pass
+class E: pass
+
+class Base(TypedDict):
+    a: A
+    b: B
+    c: C
+
+class Test(Base, total=False):
+    d: D
+    e: E
+
+class AAndB(A, B): pass
+
+test: Test
+good_keys: Literal["a", "b"]
+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]'
+reveal_type(test.get(good_keys, 3))               # N: Revealed type is 'Union[__main__.A, builtins.int, __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, builtins.int]'
+reveal_type(test.setdefault(good_keys, AAndB()))  # N: Revealed type is 'Union[__main__.A, __main__.B]'
+del test[optional_keys]
+
+
+test[bad_keys]                  # E: TypedDict "Test" has no key 'bad'
+test.get(bad_keys)              # E: TypedDict "Test" has no key 'bad'
+test.get(bad_keys, 3)           # E: TypedDict "Test" has no key 'bad'
+test.pop(good_keys)             # E: Key 'a' of TypedDict "Test" cannot be deleted \
+                                # E: Key 'b' of TypedDict "Test" cannot be deleted
+test.pop(bad_keys)              # E: Key 'a' of TypedDict "Test" cannot be deleted \
+                                # E: TypedDict "Test" has no key 'bad'
+test.setdefault(good_keys, 3)   # E: Argument 2 to "setdefault" of "TypedDict" has incompatible type "int"; expected "A"
+test.setdefault(bad_keys, 3 )   # E: Argument 2 to "setdefault" of "TypedDict" has incompatible type "int"; expected "A"
+del test[good_keys]             # E: Key 'a' of TypedDict "Test" cannot be deleted \
+                                # E: Key 'b' of TypedDict "Test" cannot be deleted
+del test[bad_keys]              # E: Key 'a' of TypedDict "Test" cannot be deleted \
+                                # E: TypedDict "Test" has no key 'bad'
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+[out]
+
 [case testLiteralIntelligentIndexingTypedDictPython2-skip]
 # flags: --python-version 2.7
 from normal_mod import NormalDict
@@ -2269,6 +2351,43 @@ UnicodeDict = TypedDict(b'UnicodeDict', {'key': int})
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-full.pyi]
 
+[case testLiteralIntelligentIndexingMultiTypedDict]
+from typing import Union
+from typing_extensions import Literal
+from mypy_extensions import TypedDict
+
+class A: pass
+class B: pass
+class C: pass
+class D: pass
+
+class D1(TypedDict):
+    a: A
+    b: B
+    c: C
+
+class D2(TypedDict):
+    b: B
+    c: C
+    d: D
+
+x: Union[D1, D2]
+bad_keys: Literal['a', 'b', 'c', 'd']
+good_keys: Literal['b', 'c']
+
+x[bad_keys]         # E: TypedDict "D1" has no key 'd' \
+                    # E: TypedDict "D2" has no key 'a'
+x.get(bad_keys)     # E: TypedDict "D1" has no key 'd' \
+                    # E: TypedDict "D2" has no key 'a'
+x.get(bad_keys, 3)  # 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]'
+reveal_type(x.get(good_keys, 3))    # N: Revealed type is 'Union[__main__.B, builtins.int, __main__.C]'
+
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
 
 --
 -- Interactions with 'Final'


Code review comments:
diff_hunk: @@ -2607,40 +2617,59 @@ def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression)
             else:
                 return union
 
-    def _get_value(self, index: Expression) -> Optional[int]:
+    def _get_values(self, index: Expression) -> Optional[List[int]]:
path: mypy/checkexpr.py
text: @user1:
Please document this method. It is not new, but now it is sufficiently non-trivial to require a docstring with examples. Also make it non-private, essentially this is the same as `try_getting_str_literals()` but for integers instead of strings.