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 #6262.

Author: Michael0x2a

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, Iterable,
     Sequence, Iterator
@@ -2465,15 +2466,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):
@@ -2489,26 +2493,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)
@@ -2521,40 +2531,60 @@ def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression)
         else:
             return UnionType.make_simplified_union(left_type.items)
 
-    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.final_value is not None:
             typ = typ.final_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.final_value is not None:
-                typ = typ.final_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.final_value is not None:
+                    key_type = key_type.final_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:

mypy/messages.py

@@ -1143,6 +1143,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)
 

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
 
 
@@ -113,18 +113,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.final_value is not None:
-        typ = typ.final_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.final_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

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, NoneTyp, UnionType, TypedDictType,
     TypeVarType
 )
+from mypy.subtypes import is_subtype
 
 
 class DefaultPlugin(Plugin):
@@ -171,26 +172,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, NoneTyp()])
+                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(NoneTyp())
+
+        return UnionType.make_simplified_union(output_types)
     return ctx.default_return_type
 
 
@@ -228,23 +237,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
 
 
@@ -273,18 +287,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
 
 
@@ -299,15 +330,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