Revert batmobile

Author: crusaderky

array_api_compat/common/_typing.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from types import ModuleType as Namespace
 from typing import (
     TYPE_CHECKING,
@@ -105,11 +106,72 @@ def shape(self, /) -> _T_co: ...
 DTypeKind: TypeAlias = _DTypeKind | tuple[_DTypeKind, ...]
 
 
+# `__array_namespace_info__.dtypes(kind="bool")`
+class DTypesBool(TypedDict):
+    bool: DType
+
+
+# `__array_namespace_info__.dtypes(kind="signed integer")`
+class DTypesSigned(TypedDict):
+    int8: DType
+    int16: DType
+    int32: DType
+    int64: DType
+
+
+# `__array_namespace_info__.dtypes(kind="unsigned integer")`
+class DTypesUnsigned(TypedDict):
+    uint8: DType
+    uint16: DType
+    uint32: DType
+    uint64: DType
+
+
+# `__array_namespace_info__.dtypes(kind="integral")`
+class DTypesIntegral(DTypesSigned, DTypesUnsigned):
+    pass
+
+
+# `__array_namespace_info__.dtypes(kind="real floating")`
+class DTypesReal(TypedDict):
+    float32: DType
+    float64: DType
+
+
+# `__array_namespace_info__.dtypes(kind="complex floating")`
+class DTypesComplex(TypedDict):
+    complex64: DType
+    complex128: DType
+
+
+# `__array_namespace_info__.dtypes(kind="numeric")`
+class DTypesNumeric(DTypesIntegral, DTypesReal, DTypesComplex):
+    pass
+
+
+# `__array_namespace_info__.dtypes(kind=None)` (default)
+class DTypesAll(DTypesBool, DTypesNumeric):
+    pass
+
+
+# `__array_namespace_info__.dtypes(kind=?)` (fallback)
+DTypesAny: TypeAlias = Mapping[str, DType]
+
+
 __all__ = [
     "Array",
     "Capabilities",
     "DType",
     "DTypeKind",
+    "DTypesAny",
+    "DTypesAll",
+    "DTypesBool",
+    "DTypesNumeric",
+    "DTypesIntegral",
+    "DTypesSigned",
+    "DTypesUnsigned",
+    "DTypesReal",
+    "DTypesComplex",
     "DefaultDTypes",
     "Device",
     "HasShape",

array_api_compat/dask/array/_info.py

@@ -12,7 +12,7 @@
 
 from __future__ import annotations
 
-from typing import Literal, TypeAlias
+from typing import Literal, TypeAlias, overload
 
 import dask.array as da
 from numpy import bool_ as bool
@@ -34,8 +34,21 @@
 )
 
 from ...common._helpers import _DASK_DEVICE, _check_device, _dask_device
-from ...common._typing import Capabilities, DefaultDTypes, DType, DTypeKind
-
+from ...common._typing import (
+    Capabilities,
+    DefaultDTypes,
+    DType,
+    DTypeKind,
+    DTypesAll,
+    DTypesAny,
+    DTypesBool,
+    DTypesComplex,
+    DTypesIntegral,
+    DTypesNumeric,
+    DTypesReal,
+    DTypesSigned,
+    DTypesUnsigned,
+)
 Device: TypeAlias = Literal["cpu"] | _dask_device
 
 
@@ -202,9 +215,41 @@ def default_dtypes(self, /, *, device: Device | None = None) -> DefaultDTypes:
             "indexing": dtype(intp),
         }
 
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: None = None
+    ) -> DTypesAll: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["bool"]
+    ) -> DTypesBool: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["signed integer"]
+    ) -> DTypesSigned: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["unsigned integer"]
+    ) -> DTypesUnsigned: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["integral"]
+    ) -> DTypesIntegral: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["real floating"]
+    ) -> DTypesReal: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["complex floating"]
+    ) -> DTypesComplex: ...
+    @overload
+    def dtypes(
+        self, /, *, device: Device | None = None, kind: Literal["numeric"]
+    ) -> DTypesNumeric: ...
     def dtypes(
         self, /, *, device: Device | None = None, kind: DTypeKind | None = None
-    ) -> dict[str, DType]:
+    ) -> DTypesAny:
         """
         The array API data types supported by Dask.