Allow TypedDict key with literal type during construction (python#7645)

JukkaL · web-flow · commit 7bbc7af1cafb · 2019-10-07T18:41:29.000+01:00
Also refactored an operation away from `mypy.plugins.common` so that it is more cleanly reusable. Also kept an alias in the original location to avoid breaking existing plugins. Fixes python#7644.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -59,7 +59,7 @@
 from mypy.plugin import Plugin, MethodContext, MethodSigContext, FunctionContext
 from mypy.typeops import (
     tuple_fallback, make_simplified_union, true_only, false_only, erase_to_union_or_bound,
-    function_type, callable_type,
+    function_type, callable_type, try_getting_str_literals
 )
 import mypy.errorcodes as codes
 
@@ -493,11 +493,19 @@ def check_typeddict_call_with_dict(self, callee: TypedDictType,
 
         item_names = []  # List[str]
         for item_name_expr, item_arg in kwargs.items:
-            if not isinstance(item_name_expr, StrExpr):
+            literal_value = None
+            if item_name_expr:
+                key_type = self.accept(item_name_expr)
+                values = try_getting_str_literals(item_name_expr, key_type)
+                if values and len(values) == 1:
+                    literal_value = values[0]
+            if literal_value is None:
                 key_context = item_name_expr or item_arg
-                self.chk.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL, key_context)
+                self.chk.fail(message_registry.TYPEDDICT_KEY_MUST_BE_STRING_LITERAL,
+                              key_context)
                 return AnyType(TypeOfAny.from_error)
-            item_names.append(item_name_expr.value)
+            else:
+                item_names.append(literal_value)
 
         return self.check_typeddict_call_with_kwargs(
             callee, OrderedDict(zip(item_names, item_args)), context)
diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py
@@ -2,16 +2,14 @@
 
 from mypy.nodes import (
     ARG_POS, MDEF, Argument, Block, CallExpr, Expression, SYMBOL_FUNCBASE_TYPES,
-    FuncDef, PassStmt, RefExpr, SymbolTableNode, Var, StrExpr,
+    FuncDef, PassStmt, RefExpr, SymbolTableNode, Var
 )
 from mypy.plugin import ClassDefContext
 from mypy.semanal import set_callable_name
-from mypy.types import (
-    CallableType, Overloaded, Type, TypeVarDef, LiteralType, Instance, UnionType,
-    get_proper_type, get_proper_types
-)
+from mypy.types import CallableType, Overloaded, Type, TypeVarDef, get_proper_type
 from mypy.typevars import fill_typevars
 from mypy.util import get_unique_redefinition_name
+from mypy.typeops import try_getting_str_literals  # noqa: F401  # Part of public API
 
 
 def _get_decorator_bool_argument(
@@ -130,38 +128,3 @@ def add_method(
 
     info.names[name] = SymbolTableNode(MDEF, func, plugin_generated=True)
     info.defn.defs.body.append(func)
-
-
-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
-    """
-    typ = get_proper_type(typ)
-
-    if isinstance(expr, StrExpr):
-        return [expr.value]
-
-    if isinstance(typ, Instance) and typ.last_known_value is not None:
-        possible_literals = [typ.last_known_value]  # type: List[Type]
-    elif isinstance(typ, UnionType):
-        possible_literals = list(typ.items)
-    else:
-        possible_literals = [typ]
-
-    strings = []
-    for lit in get_proper_types(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/typeops.py b/mypy/typeops.py
@@ -10,10 +10,12 @@
 from mypy.types import (
     TupleType, Instance, FunctionLike, Type, CallableType, TypeVarDef, Overloaded,
     TypeVarType, TypeType, UninhabitedType, FormalArgument, UnionType, NoneType,
-    AnyType, TypeOfAny, TypeType, ProperType, get_proper_type, get_proper_types, copy_type
+    AnyType, TypeOfAny, TypeType, ProperType, LiteralType, get_proper_type, get_proper_types,
+    copy_type
 )
 from mypy.nodes import (
-    FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2,
+    FuncBase, FuncItem, OverloadedFuncDef, TypeInfo, TypeVar, ARG_STAR, ARG_STAR2, Expression,
+    StrExpr
 )
 from mypy.maptype import map_instance_to_supertype
 from mypy.expandtype import expand_type_by_instance, expand_type
@@ -417,3 +419,38 @@ def callable_type(fdef: FuncItem, fallback: Instance,
         column=fdef.column,
         implicit=True,
     )
+
+
+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
+    """
+    typ = get_proper_type(typ)
+
+    if isinstance(expr, StrExpr):
+        return [expr.value]
+
+    if isinstance(typ, Instance) and typ.last_known_value is not None:
+        possible_literals = [typ.last_known_value]  # type: List[Type]
+    elif isinstance(typ, UnionType):
+        possible_literals = list(typ.items)
+    else:
+        possible_literals = [typ]
+
+    strings = []
+    for lit in get_proper_types(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/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -1883,3 +1883,25 @@ assert isinstance(u2, Mapping)
 reveal_type(u2)  # N: Revealed type is 'TypedDict('__main__.User', {'id': builtins.int, 'name': builtins.str})'
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-full.pyi]
+
+[case testTypedDictLiteralTypeKeyInCreation]
+from typing import TypedDict, Final, Literal
+
+class Value(TypedDict):
+    num: int
+
+num: Final = 'num'
+v: Value = {num: 5}
+v = {num: ''}  # E: Incompatible types (expression has type "str", TypedDict item "num" has type "int")
+
+bad: Final = 2
+v = {bad: 3}  # E: Expected TypedDict key to be string literal
+union: Literal['num', 'foo']
+v = {union: 2} # E: Expected TypedDict key to be string literal
+num2: Literal['num']
+v = {num2: 2}
+bad2: Literal['bad']
+v = {bad2: 2}  # E: Extra key 'bad' for TypedDict "Value"
+
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
