Merge branch 'main' into zarr-extensions

Author: brokkoli71

changes/2718.bugfix.rst

@@ -0,0 +1,3 @@
+0-dimensional arrays are now returning a scalar. Therefore, the return type of ``__getitem__`` changed
+to NDArrayLikeOrScalar. This change is to make the behavior of 0-dimensional arrays consistent with
+``numpy`` scalars.

src/zarr/api/asynchronous.py

@@ -38,6 +38,7 @@
     from collections.abc import Iterable
 
     from zarr.abc.codec import Codec
+    from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.storage import StoreLike
 
@@ -238,7 +239,7 @@ async def load(
     path: str | None = None,
     zarr_format: ZarrFormat | None = None,
     zarr_version: ZarrFormat | None = None,
-) -> NDArrayLike | dict[str, NDArrayLike]:
+) -> NDArrayLikeOrScalar | dict[str, NDArrayLikeOrScalar]:
     """Load data from an array or group into memory.
 
     Parameters

src/zarr/api/synchronous.py

@@ -27,7 +27,7 @@
         ShardsLike,
     )
     from zarr.core.array_spec import ArrayConfigLike
-    from zarr.core.buffer import NDArrayLike
+    from zarr.core.buffer import NDArrayLike, NDArrayLikeOrScalar
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import (
         JSON,
@@ -121,7 +121,7 @@ def load(
     path: str | None = None,
     zarr_format: ZarrFormat | None = None,
     zarr_version: ZarrFormat | None = None,
-) -> NDArrayLike | dict[str, NDArrayLike]:
+) -> NDArrayLikeOrScalar | dict[str, NDArrayLikeOrScalar]:
     """Load data from an array or group into memory.
 
     Parameters

src/zarr/core/array.py

@@ -35,6 +35,7 @@
 from zarr.core.buffer import (
     BufferPrototype,
     NDArrayLike,
+    NDArrayLikeOrScalar,
     NDBuffer,
     default_buffer_prototype,
 )
@@ -1256,7 +1257,7 @@ async def _get_selection(
         prototype: BufferPrototype,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         # check fields are sensible
         out_dtype = check_fields(fields, self.dtype)
 
@@ -1298,14 +1299,16 @@ async def _get_selection(
                 out_buffer,
                 drop_axes=indexer.drop_axes,
             )
+        if isinstance(indexer, BasicIndexer) and indexer.shape == ():
+            return out_buffer.as_scalar()
         return out_buffer.as_ndarray_like()
 
     async def getitem(
         self,
         selection: BasicSelection,
         *,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """
         Asynchronous function that retrieves a subset of the array's data based on the provided selection.
 
@@ -1318,7 +1321,7 @@ async def getitem(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayLikeOrScalar
             The retrieved subset of the array's data.
 
         Examples
@@ -2268,14 +2271,15 @@ def __array__(
             msg = "`copy=False` is not supported. This method always creates a copy."
             raise ValueError(msg)
 
-        arr_np = self[...]
+        arr = self[...]
+        arr_np: NDArrayLike = np.array(arr, dtype=dtype)
 
         if dtype is not None:
             arr_np = arr_np.astype(dtype)
 
         return arr_np
 
-    def __getitem__(self, selection: Selection) -> NDArrayLike:
+    def __getitem__(self, selection: Selection) -> NDArrayLikeOrScalar:
         """Retrieve data for an item or region of the array.
 
         Parameters
@@ -2286,8 +2290,8 @@ def __getitem__(self, selection: Selection) -> NDArrayLike:
 
         Returns
         -------
-        NDArrayLike
-             An array-like containing the data for the requested region.
+        NDArrayLikeOrScalar
+             An array-like or scalar containing the data for the requested region.
 
         Examples
         --------
@@ -2533,7 +2537,7 @@ def get_basic_selection(
         out: NDBuffer | None = None,
         prototype: BufferPrototype | None = None,
         fields: Fields | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """Retrieve data for an item or region of the array.
 
         Parameters
@@ -2551,8 +2555,8 @@ def get_basic_selection(
 
         Returns
         -------
-        NDArrayLike
-            An array-like containing the data for the requested region.
+        NDArrayLikeOrScalar
+            An array-like or scalar containing the data for the requested region.
 
         Examples
         --------
@@ -2753,7 +2757,7 @@ def get_orthogonal_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """Retrieve data by making a selection for each dimension of the array. For
         example, if an array has 2 dimensions, allows selecting specific rows and/or
         columns. The selection for each dimension can be either an integer (indexing a
@@ -2775,8 +2779,8 @@ def get_orthogonal_selection(
 
         Returns
         -------
-        NDArrayLike
-            An array-like containing the data for the requested selection.
+        NDArrayLikeOrScalar
+            An array-like or scalar containing the data for the requested selection.
 
         Examples
         --------
@@ -2989,7 +2993,7 @@ def get_mask_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """Retrieve a selection of individual items, by providing a Boolean array of the
         same shape as the array against which the selection is being made, where True
         values indicate a selected item.
@@ -3009,8 +3013,8 @@ def get_mask_selection(
 
         Returns
         -------
-        NDArrayLike
-            An array-like containing the data for the requested selection.
+        NDArrayLikeOrScalar
+            An array-like or scalar containing the data for the requested selection.
 
         Examples
         --------
@@ -3151,7 +3155,7 @@ def get_coordinate_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
 
@@ -3169,8 +3173,8 @@ def get_coordinate_selection(
 
         Returns
         -------
-        NDArrayLike
-            An array-like containing the data for the requested coordinate selection.
+        NDArrayLikeOrScalar
+            An array-like or scalar containing the data for the requested coordinate selection.
 
         Examples
         --------
@@ -3339,7 +3343,7 @@ def get_block_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayLikeOrScalar:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
 
@@ -3357,8 +3361,8 @@ def get_block_selection(
 
         Returns
         -------
-        NDArrayLike
-            An array-like containing the data for the requested block selection.
+        NDArrayLikeOrScalar
+            An array-like or scalar containing the data for the requested block selection.
 
         Examples
         --------

src/zarr/core/buffer/__init__.py

@@ -3,6 +3,7 @@
     Buffer,
     BufferPrototype,
     NDArrayLike,
+    NDArrayLikeOrScalar,
     NDBuffer,
     default_buffer_prototype,
 )
@@ -13,6 +14,7 @@
     "Buffer",
     "BufferPrototype",
     "NDArrayLike",
+    "NDArrayLikeOrScalar",
     "NDBuffer",
     "default_buffer_prototype",
     "numpy_buffer_prototype",

src/zarr/core/buffer/core.py

@@ -105,6 +105,10 @@ def __eq__(self, other: object) -> Self:  # type: ignore[explicit-override, over
         """
 
 
+ScalarType = int | float | complex | bytes | str | bool | np.generic
+NDArrayLikeOrScalar = ScalarType | NDArrayLike
+
+
 def check_item_key_is_1d_contiguous(key: Any) -> None:
     """Raises error if `key` isn't a 1d contiguous slice"""
     if not isinstance(key, slice):
@@ -419,6 +423,21 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         """
         ...
 
+    def as_scalar(self) -> ScalarType:
+        """Returns the buffer as a scalar value"""
+        if self._data.size != 1:
+            raise ValueError("Buffer does not contain a single scalar value")
+        item = self.as_numpy_array().item()
+        scalar: ScalarType
+
+        if np.issubdtype(self.dtype, np.datetime64):
+            unit: str = np.datetime_data(self.dtype)[0]  # Extract the unit (e.g., 'Y', 'D', etc.)
+            scalar = np.datetime64(item, unit)
+        else:
+            scalar = self.dtype.type(item)  # Regular conversion for non-datetime types
+
+        return scalar
+
     @property
     def dtype(self) -> np.dtype[Any]:
         return self._data.dtype

src/zarr/core/indexing.py

@@ -29,7 +29,7 @@
 
 if TYPE_CHECKING:
     from zarr.core.array import Array
-    from zarr.core.buffer import NDArrayLike
+    from zarr.core.buffer import NDArrayLikeOrScalar
     from zarr.core.chunk_grids import ChunkGrid
     from zarr.core.common import ChunkCoords
 
@@ -937,7 +937,7 @@ class OIndex:
     array: Array
 
     # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
-    def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLike:
+    def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLikeOrScalar:
         from zarr.core.array import Array
 
         # if input is a Zarr array, we materialize it now.
@@ -1046,7 +1046,7 @@ def __iter__(self) -> Iterator[ChunkProjection]:
 class BlockIndex:
     array: Array
 
-    def __getitem__(self, selection: BasicSelection) -> NDArrayLike:
+    def __getitem__(self, selection: BasicSelection) -> NDArrayLikeOrScalar:
         fields, new_selection = pop_fields(selection)
         new_selection = ensure_tuple(new_selection)
         new_selection = replace_lists(new_selection)
@@ -1236,7 +1236,9 @@ class VIndex:
     array: Array
 
     # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool
-    def __getitem__(self, selection: CoordinateSelection | MaskSelection | Array) -> NDArrayLike:
+    def __getitem__(
+        self, selection: CoordinateSelection | MaskSelection | Array
+    ) -> NDArrayLikeOrScalar:
         from zarr.core.array import Array
 
         # if input is a Zarr array, we materialize it now.

tests/test_api.py

@@ -32,6 +32,7 @@
     save_array,
     save_group,
 )
+from zarr.core.buffer import NDArrayLike
 from zarr.errors import MetadataValidationError
 from zarr.storage import MemoryStore
 from zarr.storage._utils import normalize_path
@@ -244,7 +245,9 @@ def test_open_with_mode_r(tmp_path: pathlib.Path) -> None:
     z2 = zarr.open(store=tmp_path, mode="r")
     assert isinstance(z2, Array)
     assert z2.fill_value == 1
-    assert (z2[:] == 1).all()
+    result = z2[:]
+    assert isinstance(result, NDArrayLike)
+    assert (result == 1).all()
     with pytest.raises(ValueError):
         z2[:] = 3
 
@@ -256,7 +259,9 @@ def test_open_with_mode_r_plus(tmp_path: pathlib.Path) -> None:
     zarr.ones(store=tmp_path, shape=(3, 3))
     z2 = zarr.open(store=tmp_path, mode="r+")
     assert isinstance(z2, Array)
-    assert (z2[:] == 1).all()
+    result = z2[:]
+    assert isinstance(result, NDArrayLike)
+    assert (result == 1).all()
     z2[:] = 3
 
 
@@ -272,7 +277,9 @@ async def test_open_with_mode_a(tmp_path: pathlib.Path) -> None:
     arr[...] = 1
     z2 = zarr.open(store=tmp_path, mode="a")
     assert isinstance(z2, Array)
-    assert (z2[:] == 1).all()
+    result = z2[:]
+    assert isinstance(result, NDArrayLike)
+    assert (result == 1).all()
     z2[:] = 3
 
 
@@ -284,7 +291,9 @@ def test_open_with_mode_w(tmp_path: pathlib.Path) -> None:
     arr[...] = 3
     z2 = zarr.open(store=tmp_path, mode="w", shape=(3, 3))
     assert isinstance(z2, Array)
-    assert not (z2[:] == 3).all()
+    result = z2[:]
+    assert isinstance(result, NDArrayLike)
+    assert not (result == 3).all()
     z2[:] = 3
 
 
@@ -1134,7 +1143,9 @@ def test_open_array_with_mode_r_plus(store: Store) -> None:
     zarr.ones(store=store, shape=(3, 3))
     z2 = zarr.open_array(store=store, mode="r+")
     assert isinstance(z2, Array)
-    assert (z2[:] == 1).all()
+    result = z2[:]
+    assert isinstance(result, NDArrayLike)
+    assert (result == 1).all()
     z2[:] = 3