From 60f1d6306ab574e43ecb81f7b0ff0c92a5e0bd5f Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sun, 24 Feb 2019 21:47:33 +0000 Subject: [PATCH 1/2] Support indexing unions containing tuples --- mypy/checkexpr.py | 28 ++++++++++++++++++++++------ test-data/unit/check-classes.test | 4 +--- test-data/unit/check-tuples.test | 24 ++++++++++++++++++++++++ test-data/unit/fixtures/tuple.pyi | 9 +++++++-- 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index a0ccfbe0c56a..94dbeea8c90b 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2012,14 +2012,18 @@ def check_method_call_by_name(self, arg_kinds: List[int], context: Context, local_errors: Optional[MessageBuilder] = None, + original_type: Optional[Type] = None ) -> Tuple[Type, Type]: """Type check a call to a named method on an object. - Return tuple (result type, inferred method type). + Return tuple (result type, inferred method type). The 'original_type' + is used for error messages. """ local_errors = local_errors or self.msg + original_type = original_type or base_type method_type = analyze_member_access(method, base_type, context, False, False, True, - local_errors, original_type=base_type, chk=self.chk, + local_errors, original_type=original_type, + chk=self.chk, in_literal_context=self.is_literal_context()) return self.check_method_call( method, base_type, method_type, args, arg_kinds, context, local_errors) @@ -2441,10 +2445,21 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type: # It's actually a type application. return self.accept(e.analyzed) left_type = self.accept(e.base) - if isinstance(left_type, TupleType) and self.chk.in_checked_function(): + return self.visit_index_with_type(left_type, e) + + def visit_index_with_type(self, left_type: Type, e: IndexExpr, + original_type: Optional[Type] = None) -> Type: + """Analyze type of an index expression for a given type of base expression. + + The 'original_type' is used for error messages (currently used for union types). + """ + index = e.index + if isinstance(left_type, UnionType): + return UnionType.make_simplified_union([self.visit_index_with_type(typ, e, left_type) + for typ in left_type.relevant_items()]) + elif isinstance(left_type, TupleType) and self.chk.in_checked_function(): # Special case for tuples. They return a more specific type when # indexed by an integer literal. - index = e.index if isinstance(index, SliceExpr): return self.visit_tuple_slice_helper(left_type, index) @@ -2452,7 +2467,7 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type: if n is not None: if n < 0: n += len(left_type.items) - if n >= 0 and 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) @@ -2466,7 +2481,8 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type: return self.visit_enum_index_expr(left_type.type_object(), e.index, e) else: result, method_type = self.check_method_call_by_name( - '__getitem__', left_type, [e.index], [ARG_POS], e) + '__getitem__', left_type, [e.index], [ARG_POS], e, + original_type=original_type) e.method_type = method_type return result diff --git a/test-data/unit/check-classes.test b/test-data/unit/check-classes.test index 68cb74737392..170ceb7e7058 100644 --- a/test-data/unit/check-classes.test +++ b/test-data/unit/check-classes.test @@ -4548,9 +4548,7 @@ class Bar(TypedDict): def foo(node: NodeType) -> int: x = node - # TODO: This is incorrect (https://github.com/python/mypy/issues/5930), but ensure that it - # doesn't crash at least - return x['x'] # E: Incompatible return value type (got "object", expected "int") + return x['x'] [builtins fixtures/isinstancelist.pyi] [out] diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test index 8c0f1d782fc8..c0dace5fe2b6 100644 --- a/test-data/unit/check-tuples.test +++ b/test-data/unit/check-tuples.test @@ -1209,3 +1209,27 @@ else: reveal_type(x) # E: Revealed type is 'Union[Tuple[__main__.B], None]' [builtins fixtures/tuple.pyi] + +[case testUnionOfTupleIndex] +from typing import Union, Tuple + +tup: Union[Tuple[int, str], Tuple[int, int, str]] +reveal_type(tup[0]) # E: Revealed type is 'builtins.int' +reveal_type(tup[1]) # E: Revealed type is 'Union[builtins.str, builtins.int]' +reveal_type(tup[2]) # E: Revealed type is 'Union[Any, builtins.str]' \ + # E: Tuple index out of range +reveal_type(tup[:]) # E: Revealed type is 'Union[Tuple[builtins.int, builtins.str], Tuple[builtins.int, builtins.int, builtins.str]]' + +[builtins fixtures/tuple.pyi] + +[case testUnionOfTupleIndexMixed] +from typing import Union, Tuple, List + +tup: Union[Tuple[int, str], List[int]] +reveal_type(tup[0]) # E: Revealed type is 'builtins.int' +reveal_type(tup[1]) # E: Revealed type is 'Union[builtins.str, builtins.int*]' +reveal_type(tup[2]) # E: Revealed type is 'Union[Any, builtins.int*]' \ + # E: Tuple index out of range +reveal_type(tup[:]) # E: Revealed type is 'Union[Tuple[builtins.int, builtins.str], builtins.list[builtins.int*]]' + +[builtins fixtures/tuple.pyi] diff --git a/test-data/unit/fixtures/tuple.pyi b/test-data/unit/fixtures/tuple.pyi index e400bdbace62..4e6c421a1615 100644 --- a/test-data/unit/fixtures/tuple.pyi +++ b/test-data/unit/fixtures/tuple.pyi @@ -1,6 +1,6 @@ # Builtins stub used in tuple-related test cases. -from typing import Iterable, Iterator, TypeVar, Generic, Sequence, Any +from typing import Iterable, Iterator, TypeVar, Generic, Sequence, Any, overload Tco = TypeVar('Tco', covariant=True) @@ -28,7 +28,12 @@ class unicode: pass T = TypeVar('T') -class list(Sequence[T], Generic[T]): pass +class list(Sequence[T], Generic[T]): + @overload + def __getitem__(self, i: int) -> T: ... + @overload + def __getitem__(self, s: slice) -> list[T]: ... + def isinstance(x: object, t: type) -> bool: pass def sum(iterable: Iterable[T], start: T = None) -> T: pass From 0a3d00fd5cff809bf4ea9102bd54491faa234ff6 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Tue, 26 Feb 2019 01:06:21 +0000 Subject: [PATCH 2/2] Address CR --- mypy/checkexpr.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 94dbeea8c90b..d0687780c45c 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2455,7 +2455,9 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr, """ index = e.index if isinstance(left_type, UnionType): - return UnionType.make_simplified_union([self.visit_index_with_type(typ, e, left_type) + original_type = original_type or left_type + return UnionType.make_simplified_union([self.visit_index_with_type(typ, e, + original_type) for typ in left_type.relevant_items()]) elif isinstance(left_type, TupleType) and self.chk.in_checked_function(): # Special case for tuples. They return a more specific type when