Allow slice syntax

This is useful for Annotated, and crucial for downstream libraries like torchtyping.

Author: Zac-HD

mypy/fastparse.py

@@ -1,4 +1,5 @@
 from mypy.util import unnamed_function
+import copy
 import re
 import sys
 import warnings
@@ -1551,22 +1552,38 @@ def visit_Bytes(self, n: Bytes) -> Type:
         contents = bytes_to_human_readable_repr(n.s)
         return RawExpressionType(contents, 'builtins.bytes', self.line, column=n.col_offset)
 
+    def visit_Index(self, n: ast3.Index) -> Type:
+        # cast for mypyc's benefit on Python 3.9
+        return self.visit(cast(Any, n).value)
+
+    def visit_Slice(self, n: ast3.Slice) -> Type:
+        return self.invalid_type(
+            n, note="did you mean to use ',' instead of ':' ?"
+        )
+
     # Subscript(expr value, slice slice, expr_context ctx)  # Python 3.8 and before
     # Subscript(expr value, expr slice, expr_context ctx)  # Python 3.9 and later
     def visit_Subscript(self, n: ast3.Subscript) -> Type:
         if sys.version_info >= (3, 9):  # Really 3.9a5 or later
             sliceval: Any = n.slice
-            if (isinstance(sliceval, ast3.Slice) or
-                (isinstance(sliceval, ast3.Tuple) and
-                 any(isinstance(x, ast3.Slice) for x in sliceval.elts))):
-                self.fail(TYPE_COMMENT_SYNTAX_ERROR, self.line, getattr(n, 'col_offset', -1))
-                return AnyType(TypeOfAny.from_error)
+        # Python 3.8 or earlier use a different AST structure for subscripts
+        elif isinstance(n.slice, ast3.Index):
+            sliceval: Any = n.slice.value
+        elif isinstance(n.slice, ast3.Slice):
+            sliceval = copy.deepcopy(n.slice)  # so we don't mutate passed AST
+            if getattr(sliceval, "col_offset", None) is None:
+                # Fix column information so that we get Python 3.9+ message order
+                sliceval.col_offset = sliceval.lower.col_offset
         else:
-            # Python 3.8 or earlier use a different AST structure for subscripts
-            if not isinstance(n.slice, Index):
-                self.fail(TYPE_COMMENT_SYNTAX_ERROR, self.line, getattr(n, 'col_offset', -1))
-                return AnyType(TypeOfAny.from_error)
-            sliceval = n.slice.value
+            assert isinstance(n.slice, ast3.ExtSlice)
+            dims = copy.deepcopy(n.slice.dims)
+            for s in dims:
+                if getattr(s, "col_offset", None) is None:
+                    if isinstance(s, ast3.Index):
+                        s.col_offset = s.value.col_offset  # type: ignore
+                    elif isinstance(s, ast3.Slice):
+                        s.col_offset = s.lower.col_offset  # type: ignore
+            sliceval = ast3.Tuple(dims, n.ctx)
 
         empty_tuple_index = False
         if isinstance(sliceval, ast3.Tuple):

test-data/unit/check-annotated.test

@@ -126,3 +126,20 @@ class Meta:
 x = Annotated[int, Meta()]
 reveal_type(x)  # N: Revealed type is "def () -> builtins.int"
 [builtins fixtures/tuple.pyi]
+
+[case testSliceAnnotated39]
+# flags: --python-version 3.9
+from typing_extensions import Annotated
+
+a: Annotated[int, 1:2]
+reveal_type(a)  # N: Revealed type is "builtins.int"
+
+[builtins fixtures/tuple.pyi]
+[case testSliceAnnotated38]
+# flags: --python-version 3.8
+from typing_extensions import Annotated
+
+a: Annotated[int, 1:2]
+reveal_type(a)  # N: Revealed type is "builtins.int"
+
+[builtins fixtures/tuple.pyi]

test-data/unit/check-errorcodes.test

@@ -873,3 +873,39 @@ lst: List[int] = []
 if lst:
     pass
 [builtins fixtures/list.pyi]
+
+[case testSliceInDict39]
+# flags: --python-version 3.9 --show-column-numbers
+from typing import Dict
+b: Dict[int, x:y]
+c: Dict[x:y]
+
+[builtins fixtures/dict.pyi]
+[out]
+main:3:14: error: Invalid type comment or annotation  [valid-type]
+main:3:14: note: did you mean to use ',' instead of ':' ?
+main:4:4: error: "dict" expects 2 type arguments, but 1 given  [type-arg]
+main:4:9: error: Invalid type comment or annotation  [valid-type]
+main:4:9: note: did you mean to use ',' instead of ':' ?
+
+[case testSliceInDict38]
+# flags: --python-version 3.8 --show-column-numbers
+from typing import Dict
+b: Dict[int, x:y]
+c: Dict[x:y]
+
+[builtins fixtures/dict.pyi]
+[out]
+main:3:14: error: Invalid type comment or annotation  [valid-type]
+main:3:14: note: did you mean to use ',' instead of ':' ?
+main:4:4: error: "dict" expects 2 type arguments, but 1 given  [type-arg]
+main:4:9: error: Invalid type comment or annotation  [valid-type]
+main:4:9: note: did you mean to use ',' instead of ':' ?
+
+
+[case testSliceInCustomTensorType]
+# syntactically mimics torchtyping.TensorType
+class TensorType: ...
+t: TensorType["batch":..., float]  # type: ignore
+reveal_type(t)  # N: Revealed type is "__main__.TensorType"
+[builtins fixtures/tuple.pyi]

test-data/unit/check-fastparse.test

@@ -321,8 +321,12 @@ x @= 1
 
 from typing import Dict
 x = None # type: Dict[x: y]
+
+[builtins fixtures/dict.pyi]
 [out]
-main:3: error: syntax error in type comment
+main:3: error: "dict" expects 2 type arguments, but 1 given
+main:3: error: Invalid type comment or annotation
+main:3: note: did you mean to use ',' instead of ':' ?
 
 [case testPrintStatementTrailingCommaFastParser_python2]
 

test-data/unit/parse.test

@@ -949,6 +949,21 @@ main:1: error: invalid syntax
 [out version>=3.10]
 main:1: error: invalid syntax. Perhaps you forgot a comma?
 
+[case testSliceInList39]
+# flags: --python-version 3.9
+x = [1, 2][1:2]
+[out]
+MypyFile:1(
+  AssignmentStmt:2(
+    NameExpr(x)
+    IndexExpr:2(
+      ListExpr:2(
+        IntExpr(1)
+        IntExpr(2))
+      SliceExpr:2(
+        IntExpr(1)
+        IntExpr(2)))))
+
 [case testDictionaryExpression]
 {}
 {1:x}