fix: Scalar indexing returns numpy scalar, not 0-d array (fixes #3741)

changes/3741.bugfix.md

@@ -0,0 +1 @@
+Scalar indexing on zarr Arrays now returns numpy scalars instead of 0-d ndarrays, matching numpy behavior. For example, `arr[0]` now returns a scalar value rather than a 0-dimensional array.

examples/custom_dtype/custom_dtype.py

@@ -220,7 +220,11 @@ def from_json_scalar(self, data: JSON, *, zarr_format: ZarrFormat) -> ml_dtypes.
 def test_custom_dtype(tmp_path: Path, zarr_format: ZarrFormat) -> None:
     # create array and write values
     z_w = zarr.create_array(
-        store=tmp_path, shape=(4,), dtype="int2", zarr_format=zarr_format, compressors=None
+        store=tmp_path,
+        shape=(4,),
+        dtype="int2",
+        zarr_format=zarr_format,
+        compressors=None,
     )
     z_w[:] = [-1, -2, 0, 1]
 

src/zarr/abc/buffer.py

@@ -1,4 +1,10 @@
-from zarr.core.buffer.core import ArrayLike, Buffer, BufferPrototype, NDArrayLike, NDBuffer
+from zarr.core.buffer.core import (
+    ArrayLike,
+    Buffer,
+    BufferPrototype,
+    NDArrayLike,
+    NDBuffer,
+)
 
 __all__ = [
     "ArrayLike",

src/zarr/abc/codec.py

@@ -2,7 +2,14 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, Protocol, TypeGuard, TypeVar, runtime_checkable
+from typing import (
+    TYPE_CHECKING,
+    Generic,
+    Protocol,
+    TypeGuard,
+    TypeVar,
+    runtime_checkable,
+)
 
 from typing_extensions import ReadOnly, TypedDict
 

src/zarr/abc/store.py

@@ -219,7 +219,11 @@ async def get(
         ...
 
     async def _get_bytes(
-        self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str,
+        *,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
         Retrieve raw bytes from the store asynchronously.
@@ -267,7 +271,11 @@ async def _get_bytes(
         return buffer.to_bytes()
 
     def _get_bytes_sync(
-        self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str = "",
+        *,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
     ) -> bytes:
         """
         Retrieve raw bytes from the store synchronously.
@@ -318,7 +326,11 @@ def _get_bytes_sync(
         return sync(self._get_bytes(key, prototype=prototype, byte_range=byte_range))
 
     async def _get_json(
-        self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str,
+        *,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
     ) -> Any:
         """
         Retrieve and parse JSON data from the store asynchronously.
@@ -368,7 +380,11 @@ async def _get_json(
         return json.loads(await self._get_bytes(key, prototype=prototype, byte_range=byte_range))
 
     def _get_json_sync(
-        self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
+        self,
+        key: str = "",
+        *,
+        prototype: BufferPrototype,
+        byte_range: ByteRequest | None = None,
     ) -> Any:
         """
         Retrieve and parse JSON data from the store synchronously.

src/zarr/api/asynchronous.py

@@ -108,7 +108,9 @@ def _infer_overwrite(mode: AccessModeLiteral) -> bool:
     return mode in _OVERWRITE_MODES
 
 
-def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
+def _get_shape_chunks(
+    a: ArrayLike | Any,
+) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
     """Helper function to get the shape and chunks from an array-like object"""
     shape = None
     chunks = None
@@ -179,7 +181,9 @@ def _handle_zarr_version_or_format(
         )
     if zarr_version is not None:
         warnings.warn(
-            "zarr_version is deprecated, use zarr_format", ZarrDeprecationWarning, stacklevel=2
+            "zarr_version is deprecated, use zarr_format",
+            ZarrDeprecationWarning,
+            stacklevel=2,
         )
         return zarr_version
     return zarr_format
@@ -386,7 +390,9 @@ async def open(
             is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array"
             if is_v3_array or zarr_format == 2:
                 return AsyncArray(
-                    store_path=store_path, metadata=_metadata_dict, config=kwargs.get("config")
+                    store_path=store_path,
+                    metadata=_metadata_dict,
+                    config=kwargs.get("config"),
                 )
         except (AssertionError, FileNotFoundError, NodeTypeValidationError):
             pass

src/zarr/api/synchronous.py

@@ -266,7 +266,12 @@ def save(
     """
     return sync(
         async_api.save(
-            store, *args, zarr_version=zarr_version, zarr_format=zarr_format, path=path, **kwargs
+            store,
+            *args,
+            zarr_version=zarr_version,
+            zarr_format=zarr_format,
+            path=path,
+            **kwargs,
         )
     )
 

src/zarr/codecs/__init__.py

@@ -83,7 +83,9 @@
 register_codec("numcodecs.fletcher32", Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32")
 register_codec("numcodecs.gzip", GZip, qualname="zarr.codecs.numcodecs.GZip")
 register_codec(
-    "numcodecs.jenkins_lookup3", JenkinsLookup3, qualname="zarr.codecs.numcodecs.JenkinsLookup3"
+    "numcodecs.jenkins_lookup3",
+    JenkinsLookup3,
+    qualname="zarr.codecs.numcodecs.JenkinsLookup3",
 )
 register_codec("numcodecs.pcodec", PCodec, qualname="zarr.codecs.numcodecs.PCodec")
 register_codec("numcodecs.packbits", PackBits, qualname="zarr.codecs.numcodecs.PackBits")

src/zarr/codecs/blosc.py

@@ -12,7 +12,12 @@
 
 from zarr.abc.codec import BytesBytesCodec
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
-from zarr.core.common import JSON, NamedRequiredConfig, parse_enum, parse_named_configuration
+from zarr.core.common import (
+    JSON,
+    NamedRequiredConfig,
+    parse_enum,
+    parse_named_configuration,
+)
 from zarr.core.dtype.common import HasItemSize
 
 if TYPE_CHECKING:

src/zarr/codecs/crc32c_.py

@@ -66,7 +66,8 @@ def _encode_sync(
         data = chunk_bytes.as_numpy_array()
         # Calculate the checksum and "cast" it to a numpy array
         checksum = np.array(
-            [google_crc32c.value(cast("typing_extensions.Buffer", data))], dtype=np.uint32
+            [google_crc32c.value(cast("typing_extensions.Buffer", data))],
+            dtype=np.uint32,
         )
         # Append the checksum (as bytes) to the data
         return chunk_spec.prototype.buffer.from_array_like(np.append(data, checksum.view("B")))

src/zarr/codecs/sharding.py

@@ -244,7 +244,9 @@ async def from_bytes(
 
     @classmethod
     def create_empty(
-        cls, chunks_per_shard: tuple[int, ...], buffer_prototype: BufferPrototype | None = None
+        cls,
+        chunks_per_shard: tuple[int, ...],
+        buffer_prototype: BufferPrototype | None = None,
     ) -> _ShardReader:
         if buffer_prototype is None:
             buffer_prototype = default_buffer_prototype()
@@ -297,7 +299,9 @@ def to_dict_vectorized(
 
 @dataclass(frozen=True)
 class ShardingCodec(
-    ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin
+    ArrayBytesCodec,
+    ArrayBytesCodecPartialDecodeMixin,
+    ArrayBytesCodecPartialEncodeMixin,
 ):
     """Sharding codec"""
 
@@ -312,7 +316,7 @@ def __init__(
         chunk_shape: ShapeLike,
         codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(),),
         index_codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(), Crc32cCodec()),
-        index_location: ShardingCodecIndexLocation | str = ShardingCodecIndexLocation.end,
+        index_location: (ShardingCodecIndexLocation | str) = ShardingCodecIndexLocation.end,
     ) -> None:
         chunk_shape_parsed = parse_shapelike(chunk_shape)
         codecs_parsed = parse_codecs(codecs)
@@ -585,7 +589,9 @@ async def _encode_partial_single(
 
         indexer = list(
             get_indexer(
-                selection, shape=shard_shape, chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape)
+                selection,
+                shape=shard_shape,
+                chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
             )
         )
 
@@ -701,7 +707,8 @@ def _shard_index_size(self, chunks_per_shard: tuple[int, ...]) -> int:
             get_pipeline_class()
             .from_codecs(self.index_codecs)
             .compute_encoded_size(
-                16 * product(chunks_per_shard), self._get_index_chunk_spec(chunks_per_shard)
+                16 * product(chunks_per_shard),
+                self._get_index_chunk_spec(chunks_per_shard),
             )
         )
 
@@ -746,7 +753,8 @@ async def _load_shard_index_maybe(
             )
         else:
             index_bytes = await byte_getter.get(
-                prototype=numpy_buffer_prototype(), byte_range=SuffixByteRequest(shard_index_size)
+                prototype=numpy_buffer_prototype(),
+                byte_range=SuffixByteRequest(shard_index_size),
             )
         if index_bytes is not None:
             return await self._decode_shard_index(index_bytes, chunks_per_shard)
@@ -760,7 +768,10 @@ async def _load_shard_index(
         ) or _ShardIndex.create_empty(chunks_per_shard)
 
     async def _load_full_shard_maybe(
-        self, byte_getter: ByteGetter, prototype: BufferPrototype, chunks_per_shard: tuple[int, ...]
+        self,
+        byte_getter: ByteGetter,
+        prototype: BufferPrototype,
+        chunks_per_shard: tuple[int, ...],
     ) -> _ShardReader | None:
         shard_bytes = await byte_getter.get(prototype=prototype)
 

src/zarr/codecs/zstd.py

@@ -64,7 +64,10 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:
-        return {"name": "zstd", "configuration": {"level": self.level, "checksum": self.checksum}}
+        return {
+            "name": "zstd",
+            "configuration": {"level": self.level, "checksum": self.checksum},
+        }
 
     @cached_property
     def _zstd_codec(self) -> Zstd: