Factor out normalize_dtype function

Author: tomwhite

cubed/storage/types.py

@@ -3,9 +3,8 @@
 from itertools import starmap
 from operator import mul
 
-import numpy as np
-
 from cubed.types import T_DType, T_RegularChunks, T_Shape
+from cubed.utils import normalize_dtype
 
 
 class ArrayMetadata:
@@ -16,7 +15,7 @@ def __init__(
         chunks: T_RegularChunks,
     ):
         self.shape = shape
-        self.dtype = np.dtype(dtype)
+        self.dtype = normalize_dtype(dtype)
         self.chunks = chunks
 
     @property

cubed/utils.py

@@ -29,7 +29,7 @@
 
 def array_memory(dtype: T_DType, shape: T_Shape) -> int:
     """Calculate the amount of memory in bytes that an array uses."""
-    return np.dtype(dtype).itemsize * prod(shape)
+    return normalize_dtype(dtype).itemsize * prod(shape)
 
 
 def chunk_memory(arr) -> int:
@@ -355,3 +355,13 @@ def normalize_shape(shape: Union[int, Tuple[int, ...], None]) -> Tuple[int, ...]
     shape = cast(Tuple[int, ...], shape)
     shape = tuple(int(s) for s in shape)
     return shape
+
+
+def normalize_dtype(dtype, device=None) -> T_DType:
+    """Normalize a `dtype` argument to the underlying backend array API dtype.
+
+    This allows dtypes to be specified as `bool`, `int`, `float`, or `complex`,
+    or a string (if the array API implementation supports it).
+    """
+
+    return np.dtype(dtype)