Allow SupportsIndex in slice expressions (#14738)

hauntsaninja · web-flow · commit fb6b8bc5a041 · 2023-03-11T11:47:15.000-08:00
Helps with #2410, as suggested by BvB93 in #2410 (comment) PEP 696 will be the real solution here, since it will allow us to make slice generic with few backward compatibility issues
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -4638,7 +4638,11 @@ def _super_arg_types(self, e: SuperExpr) -> Type | tuple[Type, Type]:
         return type_type, instance_type
 
     def visit_slice_expr(self, e: SliceExpr) -> Type:
-        expected = make_optional_type(self.named_type("builtins.int"))
+        try:
+            supports_index = self.chk.named_type("typing_extensions.SupportsIndex")
+        except KeyError:
+            supports_index = self.chk.named_type("builtins.int")  # thanks, fixture life
+        expected = make_optional_type(supports_index)
         for index in [e.begin_index, e.end_index, e.stride]:
             if index:
                 t = self.accept(index)
diff --git a/mypy/message_registry.py b/mypy/message_registry.py
@@ -82,7 +82,7 @@ def with_additional_msg(self, info: str) -> ErrorMessage:
 INCOMPATIBLE_TYPES_IN_CAPTURE: Final = ErrorMessage("Incompatible types in capture pattern")
 MUST_HAVE_NONE_RETURN_TYPE: Final = ErrorMessage('The return type of "{}" must be None')
 TUPLE_INDEX_OUT_OF_RANGE: Final = ErrorMessage("Tuple index out of range")
-INVALID_SLICE_INDEX: Final = ErrorMessage("Slice index must be an integer or None")
+INVALID_SLICE_INDEX: Final = ErrorMessage("Slice index must be an integer, SupportsIndex or None")
 CANNOT_INFER_LAMBDA_TYPE: Final = ErrorMessage("Cannot infer type of lambda")
 CANNOT_ACCESS_INIT: Final = (
     'Accessing "__init__" on an instance is unsound, since instance.__init__ could be from'
diff --git a/test-data/unit/check-expressions.test b/test-data/unit/check-expressions.test
@@ -1115,11 +1115,28 @@ o[:]   # E: Value of type "object" is not indexable
 
 [case testNonIntSliceBounds]
 from typing import Any
-a, o = None, None # type: (Any, object)
-a[o:1] # E: Slice index must be an integer or None
-a[1:o] # E: Slice index must be an integer or None
-a[o:]  # E: Slice index must be an integer or None
-a[:o]  # E: Slice index must be an integer or None
+a: Any
+o: object
+a[o:1] # E: Slice index must be an integer, SupportsIndex or None
+a[1:o] # E: Slice index must be an integer, SupportsIndex or None
+a[o:]  # E: Slice index must be an integer, SupportsIndex or None
+a[:o]  # E: Slice index must be an integer, SupportsIndex or None
+[builtins fixtures/slice.pyi]
+
+[case testSliceSupportsIndex]
+import typing_extensions
+class Index:
+    def __init__(self, value: int) -> None:
+        self.value = value
+    def __index__(self) -> int:
+        return self.value
+
+c = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+reveal_type(c[Index(0):Index(5)])  # N: Revealed type is "builtins.list[builtins.int]"
+[file typing_extensions.pyi]
+from typing import Protocol
+class SupportsIndex(Protocol):
+    def __index__(self) -> int: ...
 [builtins fixtures/slice.pyi]
 
 [case testNoneSliceBounds]
diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test
@@ -1248,7 +1248,7 @@ t = (0, "")
 x = 0
 y = ""
 reveal_type(t[x:])  # N: Revealed type is "builtins.tuple[Union[builtins.int, builtins.str], ...]"
-t[y:]  # E: Slice index must be an integer or None
+t[y:]  # E: Slice index must be an integer, SupportsIndex or None
 [builtins fixtures/tuple.pyi]
 
 [case testInferTupleTypeFallbackAgainstInstance]
diff --git a/test-data/unit/fixtures/slice.pyi b/test-data/unit/fixtures/slice.pyi
@@ -15,3 +15,5 @@ class str: pass
 class slice: pass
 class ellipsis: pass
 class dict: pass
+class list(Generic[T]):
+    def __getitem__(self, x: slice) -> list[T]: pass
