Allow slice syntax in e.g. Annotated[int, 1:3] and TensorType["batch":..., float] (#11345)

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
@@ -130,8 +131,6 @@ def ast3_parse(source: Union[str, bytes], filename: str, mode: str,
 
 INVALID_TYPE_IGNORE: Final = 'Invalid "type: ignore" comment'
 
-INVALID_SLICE_ERROR: Final = 'Slice usage in type annotation is invalid'
-
 TYPE_IGNORE_PATTERN: Final = re.compile(r'[^#]*#\s*type:\s*ignore\s*(.*)')
 
 
@@ -1554,22 +1553,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(INVALID_SLICE_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(INVALID_SLICE_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

@@ -897,3 +897,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

@@ -317,13 +317,6 @@ x = None  # type: Any
 x @ 1
 x @= 1
 
-[case testIncorrectTypeCommentIndex]
-
-from typing import Dict
-x = None # type: Dict[x: y]
-[out]
-main:3: error: Slice usage in type annotation is invalid
-
 [case testPrintStatementTrailingCommaFastParser_python2]
 
 print 0,

test-data/unit/parse.test

@@ -949,73 +949,6 @@ main:1: error: invalid syntax
 [out version>=3.10]
 main:1: error: invalid syntax. Perhaps you forgot a comma?
 
-[case testSliceInAnnotation39]
-# flags: --python-version 3.9
-a: Annotated[int, 1:2]  # E: Slice usage in type annotation is invalid
-b: Dict[int, x:y]  # E: Slice usage in type annotation is invalid
-c: Dict[x:y]  # E: Slice usage in type annotation is invalid
-[out]
-
-[case testSliceInAnnotation38]
-# flags: --python-version 3.8
-a: Annotated[int, 1:2]  # E: Slice usage in type annotation is invalid
-b: Dict[int, x:y]  # E: Slice usage in type annotation is invalid
-c: Dict[x:y]  # E: Slice usage in type annotation is invalid
-[out]
-
-[case testSliceInAnnotationTypeComment39]
-# flags: --python-version 3.9
-a = None  # type: Annotated[int, 1:2]  # E: Slice usage in type annotation is invalid
-b = None  # type: Dict[int, x:y]  # E: Slice usage in type annotation is invalid
-c = None  # type: Dict[x:y]  # E: Slice usage in type annotation is invalid
-[out]
-
-[case testCorrectSlicesInAnnotations39]
-# flags: --python-version 3.9
-a: Annotated[int, slice(1, 2)]
-b: Dict[int, {x:y}]
-c: Dict[{x:y}]
-[out]
-MypyFile:1(
-  AssignmentStmt:2(
-    NameExpr(a)
-    TempNode:2(
-      Any)
-    Annotated?[int?, None])
-  AssignmentStmt:3(
-    NameExpr(b)
-    TempNode:3(
-      Any)
-    Dict?[int?, None])
-  AssignmentStmt:4(
-    NameExpr(c)
-    TempNode:4(
-      Any)
-    Dict?[None]))
-
-[case testCorrectSlicesInAnnotations38]
-# flags: --python-version 3.8
-a: Annotated[int, slice(1, 2)]
-b: Dict[int, {x:y}]
-c: Dict[{x:y}]
-[out]
-MypyFile:1(
-  AssignmentStmt:2(
-    NameExpr(a)
-    TempNode:2(
-      Any)
-    Annotated?[int?, None])
-  AssignmentStmt:3(
-    NameExpr(b)
-    TempNode:3(
-      Any)
-    Dict?[int?, None])
-  AssignmentStmt:4(
-    NameExpr(c)
-    TempNode:4(
-      Any)
-    Dict?[None]))
-
 [case testSliceInList39]
 # flags: --python-version 3.9
 x = [1, 2][1:2]