diff --git a/changes/4189.bugfix.md b/changes/4189.bugfix.md new file mode 100644 index 0000000000..76ef7e7a5e --- /dev/null +++ b/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. diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 65f7767a29..548f2141d2 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -10,7 +10,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 @@ -46,6 +45,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 @@ -1225,7 +1225,7 @@ async def require_array( name: str, *, shape: ShapeLike, - dtype: npt.DTypeLike | None = None, + dtype: ZDTypeLike | None = None, exact: bool = False, **kwargs: Any, ) -> AnyAsyncArray: @@ -1239,8 +1239,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. @@ -1258,7 +1259,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})") diff --git a/tests/test_group.py b/tests/test_group.py index 29377a5392..f7f2333ef5 100644 --- a/tests/test_group.py +++ b/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 Float64, Int32 from zarr.core.dtype.common import unpack_dtype_json from zarr.core.dtype.npy.int import UInt8 from zarr.core.group import ( @@ -61,6 +62,7 @@ from zarr.core.buffer.core import Buffer from zarr.core.common import JSON, ZarrFormat + from zarr.core.dtype import ZDType, ZDTypeLike @pytest.fixture(params=["local", "memory", "zip"]) @@ -1439,6 +1441,29 @@ async def test_require_array(store: Store, zarr_format: ZarrFormat) -> None: await root.require_array("bar", shape=(10,), dtype="int8") +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + (Int32(), Int32()), + (np.dtype("int32"), Int32()), + ("int32", Int32()), + (None, Float64()), + ], + ids=["zdtype", "numpy", "str", "none"], +) +async def test_require_array_zdtype( + store: Store, zarr_format: ZarrFormat, dtype: ZDTypeLike | None, expected: ZDType[Any, Any] +) -> None: + """An existing array can be required with a ZDType, as well as a string, a NumPy dtype, + or None. See https://github.com/zarr-developers/zarr-python/issues/3377 + """ + root = await AsyncGroup.from_store(store=store, zarr_format=zarr_format) + await root.create_array("foo", shape=(10,), dtype=expected) + + foo = await root.require_array("foo", shape=(10,), dtype=dtype, exact=True) + assert foo._zdtype == expected + + @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)