MAINT: add a _canonic_name field to DTypes, refactor _info.dtypes()

Author: ev-br

array_api_strict/_dtypes.py

@@ -11,9 +11,11 @@
 
 class DType:
     _np_dtype: Final[np.dtype[Any]]
-    __slots__ = ("_np_dtype", "__weakref__")
+    _canonic_name: Final[Any]
+    __slots__ = ("_np_dtype", "_canonic_name", "__weakref__")
 
     def __init__(self, np_dtype: npt.DTypeLike):
+        self._canonic_name = np_dtype
         self._np_dtype = np.dtype(np_dtype)
 
     def __repr__(self) -> str:

array_api_strict/_info.py

@@ -6,6 +6,26 @@
 from ._typing import Capabilities, DataTypes, DefaultDataTypes
 
 
+def _kind_to_dtypes(kind : str | None = None):
+    if kind is None:
+        return {x._canonic_name:x for x in dt._all_dtypes}
+    if kind == "bool":
+        return {"bool": dt.bool}
+    if kind == "signed integer":
+        return {x._canonic_name:x for x in dt._signed_integer_dtypes}
+    if kind == "unsigned integer":
+        return {x._canonic_name:x for x in dt._unsigned_integer_dtypes}
+    if kind == "integral":
+        return {x._canonic_name:x for x in dt._integer_dtypes}
+    if kind == "real floating":
+        return {x._canonic_name:x for x in dt._real_floating_dtypes}
+    if kind == "complex floating":
+        return {x._canonic_name:x for x in dt._complex_floating_dtypes}
+    if kind == "numeric":
+        return {x._canonic_name:x for x in dt._numeric_dtypes}
+    raise ValueError(f"unsupported kind: {kind!r}")
+
+
 @requires_api_version('2023.12')
 class __array_namespace_info__:
     @requires_api_version('2023.12')
@@ -54,75 +74,9 @@ def dtypes(
         device: Device | None = None,
         kind: str | tuple[str, ...] | None = None,
     ) -> DataTypes:
-        if kind is None:
-            return {
-                "bool": dt.bool,
-                "int8": dt.int8,
-                "int16": dt.int16,
-                "int32": dt.int32,
-                "int64": dt.int64,
-                "uint8": dt.uint8,
-                "uint16": dt.uint16,
-                "uint32": dt.uint32,
-                "uint64": dt.uint64,
-                "float32": dt.float32,
-                "float64": dt.float64,
-                "complex64": dt.complex64,
-                "complex128": dt.complex128,
-            }
-        if kind == "bool":
-            return {"bool": dt.bool}
-        if kind == "signed integer":
-            return {
-                "int8": dt.int8,
-                "int16": dt.int16,
-                "int32": dt.int32,
-                "int64": dt.int64,
-            }
-        if kind == "unsigned integer":
-            return {
-                "uint8": dt.uint8,
-                "uint16": dt.uint16,
-                "uint32": dt.uint32,
-                "uint64": dt.uint64,
-            }
-        if kind == "integral":
-            return {
-                "int8": dt.int8,
-                "int16": dt.int16,
-                "int32": dt.int32,
-                "int64": dt.int64,
-                "uint8": dt.uint8,
-                "uint16": dt.uint16,
-                "uint32": dt.uint32,
-                "uint64": dt.uint64,
-            }
-        if kind == "real floating":
-            return {
-                "float32": dt.float32,
-                "float64": dt.float64,
-            }
-        if kind == "complex floating":
-            return {
-                "complex64": dt.complex64,
-                "complex128": dt.complex128,
-            }
-        if kind == "numeric":
-            return {
-                "int8": dt.int8,
-                "int16": dt.int16,
-                "int32": dt.int32,
-                "int64": dt.int64,
-                "uint8": dt.uint8,
-                "uint16": dt.uint16,
-                "uint32": dt.uint32,
-                "uint64": dt.uint64,
-                "float32": dt.float32,
-                "float64": dt.float64,
-                "complex64": dt.complex64,
-                "complex128": dt.complex128,
-            }
-        if isinstance(kind, tuple):
+        if isinstance(kind, type(None) | str):
+            return _kind_to_dtypes(kind)
+        elif isinstance(kind, tuple):
             res: DataTypes = {}
             for k in kind:
                 res.update(self.dtypes(kind=k))