fix: allow require_array to accept a ZDType

changes/4189.bugfix.md

@@ -0,0 +1 @@
+Allow `Group.require_array` to accept a `ZDType` for `dtype`, matching the other array creation methods. Previously an existing array could only be required with a string or NumPy dtype.

src/zarr/core/group.py

@@ -11,7 +11,6 @@
 from typing import TYPE_CHECKING, Literal, assert_never, cast, overload
 
 import numpy as np
-import numpy.typing as npt
 
 import zarr.api.asynchronous as async_api
 from zarr.abc.metadata import Metadata
@@ -47,6 +46,7 @@
     parse_shapelike,
 )
 from zarr.core.config import config
+from zarr.core.dtype import parse_data_type
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.metadata.io import save_metadata
 from zarr.core.sync import SyncMixin, sync
@@ -1159,7 +1159,7 @@ async def require_array(
         name: str,
         *,
         shape: ShapeLike,
-        dtype: npt.DTypeLike | None = None,
+        dtype: ZDTypeLike | None = None,
         exact: bool = False,
         **kwargs: Any,
     ) -> AnyAsyncArray:
@@ -1173,8 +1173,9 @@ async def require_array(
             Array name.
         shape : int or tuple of ints
             Array shape.
-        dtype : str or dtype, optional
-            NumPy dtype.
+        dtype : ZDTypeLike, optional
+            The data type of the array, given as a string, a NumPy dtype, or a
+            Zarr data type.
         exact : bool, optional
             If True, require `dtype` to match exactly. If false, require
             `dtype` can be cast from array dtype.
@@ -1192,7 +1193,11 @@ async def require_array(
             if shape != ds.shape:
                 raise TypeError(f"Incompatible shape ({ds.shape} vs {shape})")
 
-            dtype = np.dtype(dtype)
+            # `np.dtype(None)` used to resolve to float64 here; keep that default.
+            dtype = parse_data_type(
+                "float64" if dtype is None else dtype,
+                zarr_format=self.metadata.zarr_format,
+            ).to_native_dtype()
             if exact:
                 if ds.dtype != dtype:
                     raise TypeError(f"Incompatible dtype ({ds.dtype} vs {dtype})")

tests/test_group.py

@@ -24,6 +24,7 @@
 from zarr.core._info import GroupInfo
 from zarr.core.buffer import default_buffer_prototype
 from zarr.core.config import config as zarr_config
+from zarr.core.dtype import parse_data_type
 from zarr.core.dtype.common import unpack_dtype_json
 from zarr.core.dtype.npy.int import UInt8
 from zarr.core.group import (
@@ -1366,6 +1367,19 @@ async def test_require_array(store: Store, zarr_format: ZarrFormat) -> None:
         await root.require_array("bar", shape=(10,), dtype="int8")
 
 
+async def test_require_array_zdtype(store: Store, zarr_format: ZarrFormat) -> None:
+    """An existing array can be required with a ZDType, as well as a string or
+    a NumPy dtype. See https://github.com/zarr-developers/zarr-python/issues/3377
+    """
+    root = await AsyncGroup.from_store(store=store, zarr_format=zarr_format)
+    await root.require_array("foo", shape=(10,), dtype="int32")
+
+    zdtype = parse_data_type("int32", zarr_format=zarr_format)
+    for dtype in (zdtype, np.dtype("int32"), "int32"):
+        foo = await root.require_array("foo", shape=(10,), dtype=dtype, exact=True)
+        assert foo.dtype == np.dtype("int32")
+
+
 @pytest.mark.parametrize("consolidate", [True, False])
 async def test_members_name(store: Store, consolidate: bool, zarr_format: ZarrFormat):
     group = Group.from_store(store=store, zarr_format=zarr_format)