TypedDict: Recognize __getitem__ and __setitem__ access. (#2526)

davidfstr · JukkaL · commit 8453dda84ada · 2016-12-20T17:22:53.000Z
* TypedDict: Recognize __getitem__ and __setitem__ access.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -25,10 +25,11 @@
     DictionaryComprehension, ComplexExpr, EllipsisExpr, TypeAliasExpr,
     RefExpr, YieldExpr, BackquoteExpr, ImportFrom, ImportAll, ImportBase,
     AwaitExpr,
+    ARG_POS,
     CONTRAVARIANT, COVARIANT)
 from mypy import nodes
 from mypy.types import (
-    Type, AnyType, CallableType, Void, FunctionLike, Overloaded, TupleType,
+    Type, AnyType, CallableType, Void, FunctionLike, Overloaded, TupleType, TypedDictType,
     Instance, NoneTyp, ErrorType, strip_type, TypeType,
     UnionType, TypeVarId, TypeVarType, PartialType, DeletedType, UninhabitedType,
     true_only, false_only, function_type
@@ -1546,8 +1547,18 @@ def check_indexed_assignment(self, lvalue: IndexExpr,
         """
         self.try_infer_partial_type_from_indexed_assignment(lvalue, rvalue)
         basetype = self.accept(lvalue.base)
-        method_type = self.expr_checker.analyze_external_member_access(
-            '__setitem__', basetype, context)
+        if isinstance(basetype, TypedDictType):
+            item_type = self.expr_checker.visit_typeddict_index_expr(basetype, lvalue.index)
+            method_type = CallableType(
+                arg_types=[self.named_type('builtins.str'), item_type],
+                arg_kinds=[ARG_POS, ARG_POS],
+                arg_names=[None, None],
+                ret_type=NoneTyp(),
+                fallback=self.named_type('builtins.function')
+            )  # type: Type
+        else:
+            method_type = self.expr_checker.analyze_external_member_access(
+                '__setitem__', basetype, context)
         lvalue.method_type = method_type
         self.expr_checker.check_call(method_type, [lvalue.index, rvalue],
                                      [nodes.ARG_POS, nodes.ARG_POS],
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -1435,6 +1435,8 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type:
             else:
                 self.chk.fail(messages.TUPLE_INDEX_MUST_BE_AN_INT_LITERAL, e)
                 return AnyType()
+        elif isinstance(left_type, TypedDictType):
+            return self.visit_typeddict_index_expr(left_type, e.index)
         else:
             result, method_type = self.check_op('__getitem__', left_type, e.index, e)
             e.method_type = method_type
@@ -1481,6 +1483,18 @@ def _get_value(self, index: Expression) -> Optional[int]:
                     return -1 * operand.value
         return None
 
+    def visit_typeddict_index_expr(self, td_type: TypedDictType, index: Expression):
+        if not isinstance(index, (StrExpr, UnicodeExpr)):
+            self.msg.typeddict_item_name_must_be_string_literal(td_type, index)
+            return AnyType()
+        item_name = index.value
+
+        item_type = td_type.items.get(item_name)
+        if item_type is None:
+            self.msg.typeddict_item_name_not_found(td_type, item_name, index)
+            return AnyType()
+        return item_type
+
     def visit_cast_expr(self, expr: CastExpr) -> Type:
         """Type check a cast expression."""
         source_type = self.accept(expr.expr, context=AnyType())
diff --git a/mypy/checkmember.py b/mypy/checkmember.py
@@ -7,9 +7,11 @@
     Overloaded, TypeVarType, UnionType, PartialType,
     DeletedType, NoneTyp, TypeType, function_type
 )
-from mypy.nodes import TypeInfo, FuncBase, Var, FuncDef, SymbolNode, Context, MypyFile, TypeVarExpr
-from mypy.nodes import ARG_POS, ARG_STAR, ARG_STAR2
-from mypy.nodes import Decorator, OverloadedFuncDef
+from mypy.nodes import (
+    TypeInfo, FuncBase, Var, FuncDef, SymbolNode, Context, MypyFile, TypeVarExpr,
+    ARG_POS, ARG_STAR, ARG_STAR2,
+    Decorator, OverloadedFuncDef,
+)
 from mypy.messages import MessageBuilder
 from mypy.maptype import map_instance_to_supertype
 from mypy.expandtype import expand_type_by_instance, expand_type
diff --git a/mypy/messages.py b/mypy/messages.py
@@ -819,6 +819,19 @@ def typeddict_instantiated_with_unexpected_items(self,
         self.fail('Expected items {} but found {}.'.format(
             expected_item_names, actual_item_names), context)
 
+    def typeddict_item_name_must_be_string_literal(self,
+                                                   typ: TypedDictType,
+                                                   context: Context):
+        self.fail('Cannot prove expression is a valid item name; expected one of {}'.format(
+            format_item_name_list(typ.items.keys())), context)
+
+    def typeddict_item_name_not_found(self,
+                                      typ: TypedDictType,
+                                      item_name: str,
+                                      context: Context):
+        self.fail('\'{}\' is not a valid item name; expected one of {}'.format(
+            item_name, format_item_name_list(typ.items.keys())), context)
+
 
 def capitalize(s: str) -> str:
     """Capitalize the first character of a string."""
@@ -862,6 +875,14 @@ def format_string_list(s: Iterable[str]) -> str:
         return '%s, ... and %s (%i methods suppressed)' % (', '.join(l[:2]), l[-1], len(l) - 3)
 
 
+def format_item_name_list(s: Iterable[str]) -> str:
+    l = list(s)
+    if len(l) <= 5:
+        return '[' + ', '.join(["'%s'" % name for name in l]) + ']'
+    else:
+        return '[' + ', '.join(["'%s'" % name for name in l[:5]]) + ', ...]'
+
+
 def callable_name(type: CallableType) -> str:
     if type.name:
         return type.name
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -360,64 +360,73 @@ reveal_type(f(g))  # E: Revealed type is '<uninhabited>'
 
 -- Special Method: __getitem__
 
--- TODO: Implement support for this case.
---[case testCanGetItemOfTypedDictWithValidStringLiteralKey]
---from mypy_extensions import TypedDict
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def get_x(p: TaggedPoint) -> int:
---    return p['x']
---[builtins fixtures/dict.pyi]
+[case testCanGetItemOfTypedDictWithValidStringLiteralKey]
+from mypy_extensions import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+reveal_type(p['type'])  # E: Revealed type is 'builtins.str'
+reveal_type(p['x'])     # E: Revealed type is 'builtins.int'
+reveal_type(p['y'])     # E: Revealed type is 'builtins.int'
+[builtins fixtures/dict.pyi]
 
--- TODO: Implement support for this case.
---[case testCannotGetItemOfTypedDictWithInvalidStringLiteralKey]
---from mypy_extensions import TypedDict
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def get_z(p: TaggedPoint) -> int:
---    return p['z']  # E: ... 'z' is not a valid key for Point. Expected one of {'x', 'y'}.
---[builtins fixtures/dict.pyi]
+[case testCanGetItemOfTypedDictWithValidBytesOrUnicodeLiteralKey]
+# flags: --python-version 2.7
+from mypy_extensions import TypedDict
+Cell = TypedDict('Cell', {'value': int})
+c = Cell(value=42)
+reveal_type(c['value'])   # E: Revealed type is 'builtins.int'
+reveal_type(c[u'value'])  # E: Revealed type is 'builtins.int'
+[builtins_py2 fixtures/dict.pyi]
 
--- TODO: Implement support for this case.
---[case testCannotGetItemOfTypedDictWithNonLiteralKey]
---from mypy_extensions import TypedDict
---from typing import Union
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def get_coordinate(p: TaggedPoint, key: str) -> Union[str, int]:
---    return p[key]  # E: ... Cannot prove 'key' is a valid key for Point. Expected one of {'x', 'y'}
---[builtins fixtures/dict.pyi]
+[case testCannotGetItemOfTypedDictWithInvalidStringLiteralKey]
+from mypy_extensions import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+p['z']  # E: 'z' is not a valid item name; expected one of ['type', 'x', 'y']
+[builtins fixtures/dict.pyi]
+
+[case testCannotGetItemOfTypedDictWithNonLiteralKey]
+from mypy_extensions import TypedDict
+from typing import Union
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+def get_coordinate(p: TaggedPoint, key: str) -> Union[str, int]:
+    return p[key]  # E: Cannot prove expression is a valid item name; expected one of ['type', 'x', 'y']
+[builtins fixtures/dict.pyi]
 
 
 -- Special Method: __setitem__
 
--- TODO: Implement support for this case.
---[case testCanSetItemOfTypedDictWithValidStringLiteralKey]
---from mypy_extensions import TypedDict
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def set_x(p: TaggedPoint, x: int) -> None:
---    p['x'] = x
---[builtins fixtures/dict.pyi]
+[case testCanSetItemOfTypedDictWithValidStringLiteralKeyAndCompatibleValueType]
+from mypy_extensions import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+p['type'] = 'two_d'
+p['x'] = 1
+[builtins fixtures/dict.pyi]
 
--- TODO: Implement support for this case.
---[case testCannotSetItemOfTypedDictWithInvalidStringLiteralKey]
---from mypy_extensions import TypedDict
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def set_z(p: TaggedPoint, z: int) -> None:
---    p['z'] = z  # E: ... 'z' is not a valid key for Point. Expected one of {'x', 'y'}.
---[builtins fixtures/dict.pyi]
+[case testCannotSetItemOfTypedDictWithIncompatibleValueType]
+from mypy_extensions import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+p['x'] = 'y'  # E: Argument 2 has incompatible type "str"; expected "int"
+[builtins fixtures/dict.pyi]
 
--- TODO: Implement support for this case.
---[case testCannotSetItemOfTypedDictWithNonLiteralKey]
---from mypy_extensions import TypedDict
---from typing import Union
---TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
---p = TaggedPoint(type='2d', x=42, y=1337)
---def set_coordinate(p: TaggedPoint, key: str, value: Union[str, int]) -> None:
---    p[key] = value  # E: ... Cannot prove 'key' is a valid key for Point. Expected one of {'x', 'y'}
---[builtins fixtures/dict.pyi]
+[case testCannotSetItemOfTypedDictWithInvalidStringLiteralKey]
+from mypy_extensions import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+p['z'] = 1  # E: 'z' is not a valid item name; expected one of ['type', 'x', 'y']
+[builtins fixtures/dict.pyi]
+
+[case testCannotSetItemOfTypedDictWithNonLiteralKey]
+from mypy_extensions import TypedDict
+from typing import Union
+TaggedPoint = TypedDict('TaggedPoint', {'type': str, 'x': int, 'y': int})
+p = TaggedPoint(type='2d', x=42, y=1337)
+def set_coordinate(p: TaggedPoint, key: str, value: int) -> None:
+    p[key] = value  # E: Cannot prove expression is a valid item name; expected one of ['type', 'x', 'y']
+[builtins fixtures/dict.pyi]
 
 
 -- Special Method: get
