Allow nonliteral tuple indexing (#3514)

* Allow nonliteral tuple indexing

* fix indent

Author: ddfisher

mypy/checkexpr.py

@@ -1570,31 +1570,23 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type:
             return self.accept(e.analyzed)
         left_type = self.accept(e.base)
         if isinstance(left_type, TupleType) and self.chk.in_checked_function():
-            # Special case for tuples. They support indexing only by integer
-            # literals.
+            # 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)
 
-            ok = False
-            if isinstance(index, IntExpr):
-                n = index.value
-                ok = True
-            elif isinstance(index, UnaryExpr):
-                if index.op == '-':
-                    operand = index.expr
-                    if isinstance(operand, IntExpr):
-                        n = len(left_type.items) - operand.value
-                        ok = True
-            if ok:
+            n = self._get_value(index)
+            if n is not None:
+                if n < 0:
+                    n += len(left_type.items)
                 if n >= 0 and n < len(left_type.items):
                     return left_type.items[n]
                 else:
                     self.chk.fail(messages.TUPLE_INDEX_OUT_OF_RANGE, e)
                     return AnyType()
             else:
-                self.chk.fail(messages.TUPLE_INDEX_MUST_BE_AN_INT_LITERAL, e)
-                return AnyType()
+                return self.nonliteral_tuple_index_helper(left_type, index)
         elif isinstance(left_type, TypedDictType):
             return self.visit_typeddict_index_expr(left_type, e.index)
         elif (isinstance(left_type, CallableType)
@@ -1613,29 +1605,31 @@ def visit_tuple_slice_helper(self, left_type: TupleType, slic: SliceExpr) -> Typ
         if slic.begin_index:
             begin = self._get_value(slic.begin_index)
             if begin is None:
-                self.chk.fail(
-                    messages.TUPLE_SLICE_MUST_BE_AN_INT_LITERAL,
-                    slic.begin_index)
-                return AnyType()
+                return self.nonliteral_tuple_index_helper(left_type, slic)
 
         if slic.end_index:
             end = self._get_value(slic.end_index)
             if end is None:
-                self.chk.fail(
-                    messages.TUPLE_SLICE_MUST_BE_AN_INT_LITERAL,
-                    slic.end_index)
-                return AnyType()
+                return self.nonliteral_tuple_index_helper(left_type, slic)
 
         if slic.stride:
             stride = self._get_value(slic.stride)
             if stride is None:
-                self.chk.fail(
-                    messages.TUPLE_SLICE_MUST_BE_AN_INT_LITERAL,
-                    slic.stride)
-                return AnyType()
+                return self.nonliteral_tuple_index_helper(left_type, slic)
 
         return left_type.slice(begin, stride, end)
 
+    def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression) -> Type:
+        index_type = self.accept(index)
+        expected_type = UnionType.make_union([self.named_type('builtins.int'),
+                                              self.named_type('builtins.slice')])
+        if not self.chk.check_subtype(index_type, expected_type, index,
+                                      messages.INVALID_TUPLE_INDEX_TYPE,
+                                      'actual type', 'expected type'):
+            return AnyType()
+        else:
+            return UnionType.make_simplified_union(left_type.items)
+
     def _get_value(self, index: Expression) -> Optional[int]:
         if isinstance(index, IntExpr):
             return index.value

mypy/messages.py

@@ -54,8 +54,7 @@
 INCOMPATIBLE_TYPES_IN_YIELD_FROM = 'Incompatible types in "yield from"'
 INCOMPATIBLE_TYPES_IN_STR_INTERPOLATION = 'Incompatible types in string interpolation'
 MUST_HAVE_NONE_RETURN_TYPE = 'The return type of "{}" must be None'
-TUPLE_INDEX_MUST_BE_AN_INT_LITERAL = 'Tuple index must be an integer literal'
-TUPLE_SLICE_MUST_BE_AN_INT_LITERAL = 'Tuple slice must be an integer literal'
+INVALID_TUPLE_INDEX_TYPE = 'Invalid tuple index type'
 TUPLE_INDEX_OUT_OF_RANGE = 'Tuple index out of range'
 NEED_ANNOTATION_FOR_VAR = 'Need type annotation for variable'
 ITERABLE_EXPECTED = 'Iterable expected'

test-data/unit/check-class-namedtuple.test

@@ -542,7 +542,7 @@ class Base(NamedTuple):
         reveal_type(self.x)  # E: Revealed type is 'builtins.int'
         self.x = 3  # E: Property "x" defined in "Base" is read-only
         self[1]  # E: Tuple index out of range
-        self[T]  # E: Tuple index must be an integer literal
+        reveal_type(self[T])  # E: Revealed type is 'builtins.int'
         return self.x
     def bad_override(self) -> int:
         return self.x
@@ -571,6 +571,7 @@ reveal_type(Child(1).good_override())  # E: Revealed type is 'builtins.int'
 reveal_type(Base(1).bad_override())  # E: Revealed type is 'builtins.int'
 reveal_type(takes_base(Base(1)))  # E: Revealed type is 'builtins.int'
 reveal_type(takes_base(Child(1)))  # E: Revealed type is 'builtins.int'
+[builtins fixtures/tuple.pyi]
 
 [case testNewNamedTupleIllegalNames]
 from typing import Callable, NamedTuple

test-data/unit/check-tuples.test

@@ -163,8 +163,8 @@ b = t1[0] # E: Incompatible types in assignment (expression has type "A", variab
 t1[2]     # E: Tuple index out of range
 t1[3]     # E: Tuple index out of range
 t2[1]     # E: Tuple index out of range
-t1[n]     # E: Tuple index must be an integer literal
-t3[n:]    # E: Tuple slice must be an integer literal
+reveal_type(t1[n])     # E: Revealed type is 'Union[__main__.A, __main__.B]'
+reveal_type(t3[n:])    # E: Revealed type is 'Union[__main__.A, __main__.B, __main__.C, __main__.D, __main__.E]'
 b = t1[(0)] # E: Incompatible types in assignment (expression has type "A", variable has type "B")
 
 a = t1[0]
@@ -925,3 +925,19 @@ f((1,))  # E: Argument 1 to "f" has incompatible type "Tuple[int]"; expected "Tu
 f(('', ''))  # E: Argument 1 to "f" has incompatible type "Tuple[str, str]"; expected "Tuple[]"
 f(0)  # E: Argument 1 to "f" has incompatible type "int"; expected "Tuple[]"
 [builtins fixtures/tuple.pyi]
+
+[case testNonliteralTupleIndex]
+t = (0, "")
+x = 0
+y = ""
+reveal_type(t[x])  # E: Revealed type is 'Union[builtins.int, builtins.str]'
+t[y]  # E: Invalid tuple index type (actual type "str", expected type "Union[int, slice]")
+[builtins fixtures/tuple.pyi]
+
+[case testNonliteralTupleSlice]
+t = (0, "")
+x = 0
+y = ""
+reveal_type(t[x:])  # E: Revealed type is 'Union[builtins.int, builtins.str]'
+t[y:]  # E: Slice index must be an integer or None
+[builtins fixtures/tuple.pyi]

test-data/unit/fixtures/tuple.pyi

@@ -15,8 +15,9 @@ class tuple(Sequence[Tco], Generic[Tco]):
     def __getitem__(self, x: int) -> Tco: pass
 class function: pass
 
-# We need int for indexing tuples.
+# We need int and slice for indexing tuples.
 class int: pass
+class slice: pass
 class bool: pass
 class str: pass # For convenience
 class unicode: pass