diff --git a/mypy/checker.py b/mypy/checker.py index d6e58f9821417..ced9a1bd83485 100644 --- 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 index 6c6f58ff185d1..74eb6f0752e92 100644 --- 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 index 76f8a39cb7737..ed9007ebc6df0 100644 --- 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 index 323229f4e4940..9060923d611a5 100644 --- 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 index fc8a8d0d672e1..424c8b2b84e01 100644 --- 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 '' -- 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