Add stubs for major protocols

Author: NeilGirdhar

src/array_api_typing/__init__.py

@@ -1,6 +1,17 @@
 """Static typing support for the array API standard."""
 
-__all__ = ["HasArrayNamespace", "__version__", "__version_tuple__"]
+__all__ = [
+    "Array",
+    "DType",
+    "Device",
+    "HasArrayNamespace",
+    "Namespace",
+    "__version__",
+    "__version_tuple__",
+]
 
-from ._namespace import HasArrayNamespace
+from ._array import Array
+from ._device import Device
+from ._dtype import DType
+from ._namespace import HasArrayNamespace, Namespace
 from ._version import version as __version__, version_tuple as __version_tuple__

src/array_api_typing/_array.py

@@ -0,0 +1,8 @@
+"""Static typing support for the array API standard."""
+
+from typing import Protocol, final
+
+
+@final
+class Array(Protocol):
+    pass

src/array_api_typing/_device.py

@@ -0,0 +1,8 @@
+"""Static typing support for the array API standard."""
+
+from typing import Protocol, final
+
+
+@final
+class Device(Protocol):
+    pass

src/array_api_typing/_dtype.py

@@ -0,0 +1,8 @@
+"""Static typing support for the array API standard."""
+
+from typing import Protocol, final
+
+
+@final
+class DType(Protocol):
+    pass

src/array_api_typing/_namespace.py

@@ -1,12 +1,47 @@
 """Static typing support for the array API standard."""
 
-__all__ = ["HasArrayNamespace"]
-
-from types import ModuleType
 from typing import Protocol, final
 from typing_extensions import TypeVar
 
-T = TypeVar("T", bound=object, default=ModuleType)  # PEP 696 default
+from ._array import Array
+from ._device import Device
+from ._dtype import DType
+from ._simple import NestedSequence, SupportsBufferProtocol
+
+A = TypeVar("A", bound=Array, default=Array)  # PEP 696 default
+
+
+class Namespace(Protocol[A]):
+    """An Array API namespace."""
+
+    def asarray(
+        self,
+        obj: (
+            Array
+            | complex
+            | NestedSequence[bool | int | float | complex]
+            | SupportsBufferProtocol
+        ),
+        /,
+        *,
+        dtype: DType | None = None,
+        device: Device | None = None,
+        copy: bool | None = None,
+        **kwargs: object,
+    ) -> A: ...
+
+    def astype(
+        self,
+        x: A,
+        dtype: DType,
+        /,
+        *,
+        copy: bool = True,
+        device: Device | None = None,
+    ) -> A: ...
+
+
+T = TypeVar("T", bound=Namespace, default=Namespace)  # PEP 696 default
 
 
 @final

src/array_api_typing/_simple.py

@@ -0,0 +1,13 @@
+"""Static typing support for the array API standard."""
+
+from typing import Any, Protocol, TypeAlias, TypeVar
+
+_T_co = TypeVar("_T_co", covariant=True)
+
+
+class NestedSequence(Protocol[_T_co]):
+    def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
+    def __len__(self, /) -> int: ...
+
+
+SupportsBufferProtocol: TypeAlias = Any