Skip to content

TypedDict: Recognize __getitem__ and __setitem__ access. #2526

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Dec 20, 2016
Merged

TypedDict: Recognize __getitem__ and __setitem__ access. #2526

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
654eaf6
TypedDict: Recognize __getitem__ and __setitem__ access.
davidfstr Dec 5, 2016
7b12582
Integrate feedback.
davidfstr Dec 14, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions mypy/checker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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],
Expand Down
14 changes: 14 additions & 0 deletions mypy/checkexpr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add test case for both 'foo' and u'foo' key names in Python 2 mode.

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())
Expand Down
8 changes: 5 additions & 3 deletions mypy/checkmember.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 21 additions & 0 deletions mypy/messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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
Expand Down
109 changes: 59 additions & 50 deletions test-data/unit/check-typeddict.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also try p['x'] = 'y' (should be an error) and p['type'] = 'a' (valid).

[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
Expand Down