diff --git a/changes/4413.bugfix.md b/changes/4413.bugfix.md new file mode 100644 index 0000000000..9a0650d1c1 --- /dev/null +++ b/changes/4413.bugfix.md @@ -0,0 +1 @@ +`Group.require_array` now applies the `config` argument when it returns an existing array. Previously `config` was only used when `require_array` created a new array, so an existing array was returned with the global defaults. An invalid `config` now raises instead of being ignored. diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 76cce0c8df..aaca4084c8 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -28,6 +28,7 @@ _parse_deprecated_compressor, create_array, ) +from zarr.core.array_spec import parse_array_config from zarr.core.attributes import Attributes from zarr.core.buffer import default_buffer_prototype from zarr.core.common import ( @@ -1240,6 +1241,7 @@ async def require_array( shape: ShapeLike, dtype: ZDTypeLike | None = None, exact: bool = False, + config: ArrayConfigLike | None = None, **kwargs: Any, ) -> AnyAsyncArray: """Obtain an array, creating if it doesn't exist. @@ -1258,6 +1260,9 @@ async def require_array( exact : bool, optional If True, require `dtype` to match exactly. If false, require `dtype` can be cast from array dtype. + config : ArrayConfigLike or None, default=None + Runtime configuration for the array, whether it is created or already exists. + Keys not specified are taken from the global configuration. Returns ------- @@ -1284,8 +1289,13 @@ async def require_array( if not np.can_cast(ds.dtype, dtype): raise TypeError(f"Incompatible dtype ({ds.dtype} vs {dtype})") except KeyError: - ds = await self.create_array(name, shape=shape, dtype=dtype, **kwargs) + return await self.create_array(name, shape=shape, dtype=dtype, config=config, **kwargs) + # `config` is the runtime configuration of the returned array, not stored metadata, + # so it applies to an existing array too. It is parsed the same way as on the create + # branch: missing keys come from the global configuration, unknown keys raise. + if config is not None: + return ds.with_config(parse_array_config(config)) return ds async def update_attributes(self, new_attributes: dict[str, Any]) -> AsyncGroup: @@ -2837,7 +2847,14 @@ def create_array( ) ) - def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray: + def require_array( + self, + name: str, + *, + shape: ShapeLike, + config: ArrayConfigLike | None = None, + **kwargs: Any, + ) -> AnyArray: """Obtain an array, creating if it doesn't exist. Other `kwargs` are as per [zarr.Group.create_array][]. @@ -2846,6 +2863,9 @@ def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArr ---------- name : str Array name. + config : ArrayConfigLike or None, default=None + Runtime configuration for the array, whether it is created or already exists. + Keys not specified are taken from the global configuration. **kwargs : See [zarr.Group.create_array][]. @@ -2853,7 +2873,9 @@ def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArr ------- a : Array """ - return Array(self._sync(self._async_group.require_array(name, shape=shape, **kwargs))) + return Array( + self._sync(self._async_group.require_array(name, shape=shape, config=config, **kwargs)) + ) def empty(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> AnyArray: """Create an empty array with the specified shape in this Group. The contents will be filled with diff --git a/tests/test_group.py b/tests/test_group.py index 31fbd138cd..a756354dca 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -22,6 +22,7 @@ from zarr.abc.store import Store from zarr.core import sync_group from zarr.core._info import GroupInfo +from zarr.core.array_spec import ArrayConfig from zarr.core.buffer import default_buffer_prototype from zarr.core.config import config as zarr_config from zarr.core.dtype import Float64, Int32 @@ -62,6 +63,7 @@ from collections.abc import Callable from zarr.core.array import ShardsLike + from zarr.core.array_spec import ArrayConfigLike from zarr.core.buffer.core import Buffer from zarr.core.common import JSON, ChunksLike, ZarrFormat from zarr.core.dtype import ZDType, ZDTypeLike @@ -1467,6 +1469,76 @@ async def test_require_array(store: Store, zarr_format: ZarrFormat) -> None: await root.require_array("bar", shape=(10,), dtype="int8") +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("api", ["sync", "async"]) +@pytest.mark.parametrize("exists", [True, False]) +@pytest.mark.parametrize( + ("config", "expected"), + [ + (None, ArrayConfig.from_dict({})), + ( + {"read_missing_chunks": False}, + ArrayConfig.from_dict({"read_missing_chunks": False}), + ), + ( + {"order": "F", "write_empty_chunks": True}, + ArrayConfig.from_dict({"order": "F", "write_empty_chunks": True}), + ), + ( + ArrayConfig.from_dict({"read_missing_chunks": False}), + ArrayConfig.from_dict({"read_missing_chunks": False}), + ), + ], +) +async def test_require_array_config( + zarr_format: ZarrFormat, + api: Literal["sync", "async"], + exists: bool, + config: ArrayConfigLike | None, + expected: ArrayConfig, +) -> None: + """ + require_array applies the config argument whether it creates the array or returns an + existing one, with keys missing from a partial config taken from the global defaults. + """ + group = Group.from_store(MemoryStore(), zarr_format=zarr_format) + if exists: + group.create_array("a", shape=(4,), dtype="uint8") + if api == "sync": + observed = group.require_array("a", shape=(4,), dtype="uint8", config=config).config + else: + observed = ( + await group._async_group.require_array("a", shape=(4,), dtype="uint8", config=config) + ).config + assert observed == expected + + +@pytest.mark.parametrize("exists", [True, False]) +def test_require_array_unknown_config_key(exists: bool) -> None: + """ + Requiring an array with a config containing an unknown key raises TypeError, whether the + array exists or is created. + """ + group = Group.from_store(MemoryStore()) + if exists: + group.create_array("a", shape=(4,), dtype="uint8") + with pytest.raises(TypeError, match=r"Unknown array config keys: \['nope'\]"): + group.require_array("a", shape=(4,), dtype="uint8", config={"nope": 1}) + + +@pytest.mark.parametrize("exists", [True, False]) +def test_require_array_invalid_config_value(exists: bool) -> None: + """ + Requiring an array with a config containing an invalid value raises ValueError, whether the + array exists or is created. + """ + group = Group.from_store(MemoryStore()) + if exists: + group.create_array("a", shape=(4,), dtype="uint8") + with pytest.raises(ValueError, match="Expected instance of bool"): + group.require_array("a", shape=(4,), dtype="uint8", config={"read_missing_chunks": "yes"}) + + @pytest.mark.parametrize( ("dtype", "expected"), [