use builtins.ellipsis for static checking

Author: pmeier

consumer.py

@@ -1,5 +1,17 @@
 import enum
-from typing import Any, Optional, Tuple, TypeVar, Union, Protocol, runtime_checkable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Optional,
+    Protocol,
+    Tuple,
+    TypeVar,
+    Union,
+    runtime_checkable,
+)
+
+if TYPE_CHECKING:
+    from builtins import ellipsis
 
 
 VendoredDtype = Any
@@ -83,7 +95,13 @@ def __ge__(self, other: Union[int, float, A], /) -> A:
 
     def __getitem__(
         self,
-        key: Union[int, slice, Tuple[Union[int, slice], ...], A],
+        key: Union[
+            int,
+            slice,
+            "ellipsis",
+            Tuple[Union[int, slice, "ellipsis"], ...],
+            A,
+        ],
         /,
     ) -> A:
         ...
@@ -142,7 +160,13 @@ def __rshift__(self, other: Union[int, A], /) -> A:
 
     def __setitem__(
         self,
-        key: Union[int, slice, Tuple[Union[int, slice], ...], A],
+        key: Union[
+            int,
+            slice,
+            "ellipsis",
+            Tuple[Union[int, slice, "ellipsis"], ...],
+            A,
+        ],
         value: Union[bool, int, float, A],
         /,
     ) -> None:

provider.py

@@ -1,5 +1,8 @@
 import enum
-from typing import Union, Optional, Any, Tuple
+from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    from builtins import ellipsis
 
 
 class ArrayImplementation:
@@ -82,7 +85,8 @@ def __getitem__(
         key: Union[
             int,
             slice,
-            Tuple[Union[int, slice], ...],
+            "ellipsis",
+            Tuple[Union[int, slice, "ellipsis"], ...],
             "ArrayImplementation",
         ],
         /,
@@ -166,7 +170,8 @@ def __setitem__(
         key: Union[
             int,
             slice,
-            Tuple[Union[int, slice], ...],
+            "ellipsis",
+            Tuple[Union[int, slice, "ellipsis"], ...],
             "ArrayImplementation",
         ],
         value: Union[bool, int, float, "ArrayImplementation"],

runtime_checks.py

@@ -1,5 +1,4 @@
 from consumer import VendoredArrayProtocol
 from provider import ArrayImplementation
 
-
 assert isinstance(ArrayImplementation(), VendoredArrayProtocol)

static_checks.py

@@ -1,16 +1,15 @@
 import enum
-from typing import TypeVar, Union, Any, Tuple, Optional
+from typing import Any, Optional, Tuple, TypeVar, Union
 
-from provider import ArrayImplementation
 from consumer import (
+    VendoredArrayNamespace,
     VendoredArrayProtocol,
-    VendoredDtype,
     VendoredDevice,
-    VendoredArrayNamespace,
+    VendoredDtype,
     VendoredPyCapsule,
     VendoredShape,
 )
-
+from provider import ArrayImplementation
 
 AIT = ArrayImplementation
 AI = AIT()