Merge branch 'main' into test-stateful-hierarchy

Author: dcherian

src/zarr/abc/store.py

@@ -284,7 +284,6 @@ async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None:
         Insert multiple (key, value) pairs into storage.
         """
         await gather(*starmap(self.set, values))
-        return
 
     @property
     @abstractmethod

src/zarr/api/asynchronous.py

@@ -10,6 +10,7 @@
 
 from zarr.abc.store import Store
 from zarr.core.array import Array, AsyncArray, get_array_metadata
+from zarr.core.buffer import NDArrayLike
 from zarr.core.common import (
     JSON,
     AccessModeLiteral,
@@ -31,7 +32,6 @@
     from collections.abc import Iterable
 
     from zarr.abc.codec import Codec
-    from zarr.core.buffer import NDArrayLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
 
     # TODO: this type could use some more thought
@@ -393,6 +393,8 @@ async def save_array(
         _handle_zarr_version_or_format(zarr_version=zarr_version, zarr_format=zarr_format)
         or _default_zarr_version()
     )
+    if not isinstance(arr, NDArrayLike):
+        raise TypeError("arr argument must be numpy or other NDArrayLike array")
 
     mode = kwargs.pop("mode", None)
     store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
@@ -447,16 +449,26 @@ async def save_group(
         or _default_zarr_version()
     )
 
+    for arg in args:
+        if not isinstance(arg, NDArrayLike):
+            raise TypeError(
+                "All arguments must be numpy or other NDArrayLike arrays (except store, path, storage_options, and zarr_format)"
+            )
+    for k, v in kwargs.items():
+        if not isinstance(v, NDArrayLike):
+            raise TypeError(f"Keyword argument '{k}' must be a numpy or other NDArrayLike array")
+
     if len(args) == 0 and len(kwargs) == 0:
         raise ValueError("at least one array must be provided")
     aws = []
     for i, arr in enumerate(args):
+        _path = f"{path}/arr_{i}" if path is not None else f"arr_{i}"
         aws.append(
             save_array(
                 store,
                 arr,
                 zarr_format=zarr_format,
-                path=f"{path}/arr_{i}",
+                path=_path,
                 storage_options=storage_options,
             )
         )

src/zarr/codecs/__init__.py

@@ -9,15 +9,13 @@
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.codecs.gzip import GzipCodec
-from zarr.codecs.pipeline import BatchedCodecPipeline
 from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
 from zarr.codecs.zstd import ZstdCodec
 from zarr.core.metadata.v3 import DataType
 
 __all__ = [
-    "BatchedCodecPipeline",
     "BloscCname",
     "BloscCodec",
     "BloscShuffle",

src/zarr/codecs/registry.py

@@ -252,7 +252,7 @@ def create_empty(
     def __setitem__(self, chunk_coords: ChunkCoords, value: Buffer) -> None:
         chunk_start = len(self.buf)
         chunk_length = len(value)
-        self.buf = self.buf + value
+        self.buf += value
         self.index.set_chunk_slice(chunk_coords, slice(chunk_start, chunk_start + chunk_length))
 
     def __delitem__(self, chunk_coords: ChunkCoords) -> None:

src/zarr/codecs/sharding.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+from zarr.core.buffer import Buffer, NDBuffer  # noqa: F401
+from zarr.core.codec_pipeline import BatchedCodecPipeline  # noqa: F401

src/zarr/core/__init__.py

@@ -47,7 +47,7 @@ def reset(self) -> None:
             "threading": {"max_workers": None},
             "json_indent": 2,
             "codec_pipeline": {
-                "path": "zarr.codecs.pipeline.BatchedCodecPipeline",
+                "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",
                 "batch_size": 1,
             },
             "codecs": {

src/zarr/codecs/pipeline.py renamed to src/zarr/core/codec_pipeline.py

@@ -675,7 +675,7 @@ def check(a: npt.NDArray[Any]) -> Order:
 def wraparound_indices(x: npt.NDArray[Any], dim_len: int) -> None:
     loc_neg = x < 0
     if np.any(loc_neg):
-        x[loc_neg] = x[loc_neg] + dim_len
+        x[loc_neg] += dim_len
 
 
 def boundscheck_indices(x: npt.NDArray[Any], dim_len: int) -> None:
@@ -1000,8 +1000,8 @@ def __init__(
                 if stop < 0:
                     stop = dim_numchunks + stop
 
-                start = start * dim_chunk_size
-                stop = stop * dim_chunk_size
+                start *= dim_chunk_size
+                stop *= dim_chunk_size
                 slice_ = slice(start, stop)
 
             else:

src/zarr/core/config.py

@@ -132,6 +132,33 @@ async def test_open_group_unspecified_version(
         assert g2.metadata.zarr_format == zarr_format
 
 
+@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
+@pytest.mark.parametrize("n_args", [10, 1, 0])
+@pytest.mark.parametrize("n_kwargs", [10, 1, 0])
+def test_save(store: Store, n_args: int, n_kwargs: int) -> None:
+    data = np.arange(10)
+    args = [np.arange(10) for _ in range(n_args)]
+    kwargs = {f"arg_{i}": data for i in range(n_kwargs)}
+
+    if n_kwargs == 0 and n_args == 0:
+        with pytest.raises(ValueError):
+            save(store)
+    elif n_args == 1 and n_kwargs == 0:
+        save(store, *args)
+        array = open(store)
+        assert isinstance(array, Array)
+        assert_array_equal(array[:], data)
+    else:
+        save(store, *args, **kwargs)  # type: ignore[arg-type]
+        group = open(store)
+        assert isinstance(group, Group)
+        for array in group.array_values():
+            assert_array_equal(array[:], data)
+        for k in kwargs.keys():
+            assert k in group
+        assert group.nmembers() == n_args + n_kwargs
+
+
 def test_save_errors() -> None:
     with pytest.raises(ValueError):
         # no arrays provided
@@ -142,6 +169,10 @@ def test_save_errors() -> None:
     with pytest.raises(ValueError):
         # no arrays provided
         save("data/group.zarr")
+    with pytest.raises(TypeError):
+        # mode is no valid argument and would get handled as an array
+        a = np.arange(10)
+        zarr.save("data/example.zarr", a, mode="w")
 
 
 def test_open_with_mode_r(tmp_path: pathlib.Path) -> None: