Repository: python/mypy

Issue:
Allow non-literals when indexing uniform tuples
Currently mypy doesn't allow indexing fixed-length tuples using a non-literal index. We should probably support this at least for tuples with homogeneous types. 

For example, code like this should probably be okay:

``` python
t = (1, 3, 7)
for i in range(len(t)):
    print(i, t[i])
```

The example below is actually already valid so rejecting the above code is inconsistent (this works due to how we set the `fallback` attribute for tuples to be `Tuple[t, ...]`):

``` python
t = (1, 2, 3)
for i, x in enumerate(t):
    print(i, x)
```



Discussion:
The relevant bit of code is `visit_index_expr_helper` in `mypy/checkexpr.py`.

I think there's a good argument for allowing it for all tuples, just with a return type of `Union[T, U, V...]`.

Rather than forcing the error immediately, let the error be deferred if necessary.

I had a generator expression that constructed a tuple, that I knew from context would be of length 3
```
self.tpl = tuple(int(x) for x in some_string.split('.'))  # type: Tuple[int, int, int]
```
And then later `__getitem__` that indexed into the tuple
```
def __getitem__(self, key):
    # type: (int) -> int
    return self.tpl[key]  
```
There were two mypy errors:
```
error: Incompatible types in assignment (expression has type Tuple[int, ...], variable has type "Tuple[int, int, int]")
error: Tuple index must be an integer literal
```
The first is easier to understand, the latter I found very confusing because it went away as soon as I typed the tuple-generator exception without a specific length, e.g. `Tuple[int, ...]`. Why would indexing into a less-specific tuple be ok but a fixed-length one bad?

Resolving patch:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
--- a/mypy/checkexpr.py
+++ b/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
diff --git a/mypy/messages.py b/mypy/messages.py
--- a/mypy/messages.py
+++ b/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 patch:
diff --git a/test-data/unit/check-class-namedtuple.test b/test-data/unit/check-class-namedtuple.test
--- a/test-data/unit/check-class-namedtuple.test
+++ b/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
diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test
--- a/test-data/unit/check-tuples.test
+++ b/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]