create NDArrayOrScalarLike

Author: brokkoli71

src/zarr/api/asynchronous.py

@@ -11,7 +11,7 @@
 
 from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
 from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
-from zarr.core.buffer import NDArrayLike
+from zarr.core.buffer import NDArrayOrScalarLike
 from zarr.core.common import (
     JSON,
     AccessModeLiteral,
@@ -232,7 +232,7 @@ async def load(
     path: str | None = None,
     zarr_format: ZarrFormat | None = None,
     zarr_version: ZarrFormat | None = None,
-) -> NDArrayLike | dict[str, NDArrayLike]:
+) -> NDArrayOrScalarLike | dict[str, NDArrayOrScalarLike]:
     """Load data from an array or group into memory.
 
     Parameters
@@ -348,7 +348,7 @@ async def open_consolidated(
 
 async def save(
     store: StoreLike,
-    *args: NDArrayLike,
+    *args: NDArrayOrScalarLike,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
     path: str | None = None,
@@ -381,7 +381,7 @@ async def save(
 
 async def save_array(
     store: StoreLike,
-    arr: NDArrayLike,
+    arr: NDArrayOrScalarLike,
     *,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
@@ -412,8 +412,8 @@ async def save_array(
         _handle_zarr_version_or_format(zarr_version=zarr_version, zarr_format=zarr_format)
         or _default_zarr_format()
     )
-    if not isinstance(arr, NDArrayLike):
-        raise TypeError("arr argument must be numpy or other NDArrayLike array")
+    if not isinstance(arr, NDArrayOrScalarLike):
+        raise TypeError("arr argument must be numpy or other NDArrayOrScalarLike array")
 
     mode = kwargs.pop("mode", "a")
     store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
@@ -436,12 +436,12 @@ async def save_array(
 
 async def save_group(
     store: StoreLike,
-    *args: NDArrayLike,
+    *args: NDArrayOrScalarLike,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
     path: str | None = None,
     storage_options: dict[str, Any] | None = None,
-    **kwargs: NDArrayLike,
+    **kwargs: NDArrayOrScalarLike,
 ) -> None:
     """Convenience function to save several NumPy arrays to the local file system, following a
     similar API to the NumPy savez()/savez_compressed() functions.
@@ -474,13 +474,15 @@ async def save_group(
     )
 
     for arg in args:
-        if not isinstance(arg, NDArrayLike):
+        if not isinstance(arg, NDArrayOrScalarLike):
             raise TypeError(
-                "All arguments must be numpy or other NDArrayLike arrays (except store, path, storage_options, and zarr_format)"
+                "All arguments must be numpy or other NDArrayOrScalarLike 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 not isinstance(v, NDArrayOrScalarLike):
+            raise TypeError(
+                f"Keyword argument '{k}' must be a numpy or other NDArrayOrScalarLike array"
+            )
 
     if len(args) == 0 and len(kwargs) == 0:
         raise ValueError("at least one array must be provided")

src/zarr/api/synchronous.py

@@ -26,7 +26,7 @@
         ShardsLike,
     )
     from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
-    from zarr.core.buffer import NDArrayLike
+    from zarr.core.buffer import NDArrayOrScalarLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import (
         JSON,
@@ -119,7 +119,7 @@ def load(
     path: str | None = None,
     zarr_format: ZarrFormat | None = None,
     zarr_version: ZarrFormat | None = None,
-) -> NDArrayLike | dict[str, NDArrayLike]:
+) -> NDArrayOrScalarLike | dict[str, NDArrayOrScalarLike]:
     """Load data from an array or group into memory.
 
     Parameters
@@ -216,7 +216,7 @@ def open_consolidated(*args: Any, use_consolidated: Literal[True] = True, **kwar
 
 def save(
     store: StoreLike,
-    *args: NDArrayLike,
+    *args: NDArrayOrScalarLike,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
     path: str | None = None,
@@ -247,7 +247,7 @@ def save(
 @_deprecate_positional_args
 def save_array(
     store: StoreLike,
-    arr: NDArrayLike,
+    arr: NDArrayOrScalarLike,
     *,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
@@ -290,12 +290,12 @@ def save_array(
 
 def save_group(
     store: StoreLike,
-    *args: NDArrayLike,
+    *args: NDArrayOrScalarLike,
     zarr_version: ZarrFormat | None = None,  # deprecated
     zarr_format: ZarrFormat | None = None,
     path: str | None = None,
     storage_options: dict[str, Any] | None = None,
-    **kwargs: NDArrayLike,
+    **kwargs: NDArrayOrScalarLike,
 ) -> None:
     """Save several NumPy arrays to the local file system.
 

src/zarr/codecs/bytes.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from zarr.abc.codec import ArrayBytesCodec
-from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
+from zarr.core.buffer import Buffer, NDArrayOrScalarLike, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.registry import register_codec
 
@@ -81,7 +81,7 @@ async def _decode_single(
             dtype = np.dtype(f"|{chunk_spec.dtype.str[1:]}")
 
         as_array_like = chunk_bytes.as_array_like()
-        if isinstance(as_array_like, NDArrayLike):
+        if isinstance(as_array_like, NDArrayOrScalarLike):
             as_nd_array_like = as_array_like
         else:
             as_nd_array_like = np.asanyarray(as_array_like)

src/zarr/core/array.py

@@ -34,7 +34,7 @@
 from zarr.core.attributes import Attributes
 from zarr.core.buffer import (
     BufferPrototype,
-    NDArrayLike,
+    NDArrayOrScalarLike,
     NDBuffer,
     default_buffer_prototype,
 )
@@ -1255,7 +1255,7 @@ async def _get_selection(
         prototype: BufferPrototype,
         out: NDBuffer | None = None,
         fields: Fields | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         # check fields are sensible
         out_dtype = check_fields(fields, self.dtype)
 
@@ -1305,7 +1305,7 @@ async def getitem(
         selection: BasicSelection,
         *,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """
         Asynchronous function that retrieves a subset of the array's data based on the provided selection.
 
@@ -1318,7 +1318,7 @@ async def getitem(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             The retrieved subset of the array's data.
 
         Examples
@@ -1396,11 +1396,11 @@ async def _set_selection(
             # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}"
             if not hasattr(value, "dtype") or value.dtype.name != self.metadata.dtype.name:
                 if hasattr(value, "astype"):
-                    # Handle things that are already NDArrayLike more efficiently
+                    # Handle things that are already NDArrayOrScalarLike more efficiently
                     value = value.astype(dtype=self.metadata.dtype, order="A")
                 else:
                     value = np.array(value, dtype=self.metadata.dtype, order="A")
-        value = cast(NDArrayLike, value)
+        value = cast(NDArrayOrScalarLike, value)
         # We accept any ndarray like object from the user and convert it
         # to a NDBuffer (or subclass). From this point onwards, we only pass
         # Buffer and NDBuffer between components.
@@ -2260,7 +2260,7 @@ def _iter_chunk_regions(
 
     def __array__(
         self, dtype: npt.DTypeLike | None = None, copy: bool | None = None
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """
         This method is used by numpy when converting zarr.Array into a numpy array.
         For more information, see https://numpy.org/devdocs/user/basics.interoperability.html#the-array-method
@@ -2278,7 +2278,7 @@ def __array__(
 
         return arr_np
 
-    def __getitem__(self, selection: Selection) -> NDArrayLike:
+    def __getitem__(self, selection: Selection) -> NDArrayOrScalarLike:
         """Retrieve data for an item or region of the array.
 
         Parameters
@@ -2289,7 +2289,7 @@ def __getitem__(self, selection: Selection) -> NDArrayLike:
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
              An array-like containing the data for the requested region.
 
         Examples
@@ -2536,7 +2536,7 @@ def get_basic_selection(
         out: NDBuffer | None = None,
         prototype: BufferPrototype | None = None,
         fields: Fields | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """Retrieve data for an item or region of the array.
 
         Parameters
@@ -2554,7 +2554,7 @@ def get_basic_selection(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             An array-like containing the data for the requested region.
 
         Examples
@@ -2756,7 +2756,7 @@ def get_orthogonal_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """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
@@ -2778,7 +2778,7 @@ def get_orthogonal_selection(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             An array-like containing the data for the requested selection.
 
         Examples
@@ -2992,7 +2992,7 @@ def get_mask_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """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.
@@ -3012,7 +3012,7 @@ def get_mask_selection(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             An array-like containing the data for the requested selection.
 
         Examples
@@ -3154,7 +3154,7 @@ def get_coordinate_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
 
@@ -3172,7 +3172,7 @@ def get_coordinate_selection(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             An array-like containing the data for the requested coordinate selection.
 
         Examples
@@ -3325,7 +3325,7 @@ def set_coordinate_selection(
             value = np.array(value).reshape(-1)
 
         if not is_scalar(value, self.dtype) and (
-            isinstance(value, NDArrayLike) and indexer.shape != value.shape
+            isinstance(value, NDArrayOrScalarLike) and indexer.shape != value.shape
         ):
             raise ValueError(
                 f"Attempting to set a selection of {indexer.sel_shape[0]} "
@@ -3342,7 +3342,7 @@ def get_block_selection(
         out: NDBuffer | None = None,
         fields: Fields | None = None,
         prototype: BufferPrototype | None = None,
-    ) -> NDArrayLike:
+    ) -> NDArrayOrScalarLike:
         """Retrieve a selection of individual items, by providing the indices
         (coordinates) for each selected item.
 
@@ -3360,7 +3360,7 @@ def get_block_selection(
 
         Returns
         -------
-        NDArrayLike
+        NDArrayOrScalarLike
             An array-like containing the data for the requested block selection.
 
         Examples

src/zarr/core/buffer/__init__.py

@@ -2,7 +2,7 @@
     ArrayLike,
     Buffer,
     BufferPrototype,
-    NDArrayLike,
+    NDArrayOrScalarLike,
     NDBuffer,
     default_buffer_prototype,
 )
@@ -12,7 +12,7 @@
     "ArrayLike",
     "Buffer",
     "BufferPrototype",
-    "NDArrayLike",
+    "NDArrayOrScalarLike",
     "NDBuffer",
     "default_buffer_prototype",
     "numpy_buffer_prototype",

src/zarr/core/buffer/core.py

@@ -9,6 +9,7 @@
     NamedTuple,
     Protocol,
     SupportsIndex,
+    Union,
     cast,
     runtime_checkable,
 )
@@ -106,6 +107,9 @@ def __eq__(self, other: object) -> Self:  # type: ignore[explicit-override, over
         """
 
 
+NDArrayLike = Union["np.ScalarType", NDArrayLike]
+
+
 class ScalarWrapper:
     def __init__(self, value: Any, dtype: np.dtype[Any] | None = None) -> None:
         self._value: Any = value
@@ -552,7 +556,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
         """
         ...
 
-    def as_scalar(self) -> np.generic:
+    def as_scalar(self) -> np.ScalarType:
         """Returns the buffer as a scalar value
 
         Returns

src/zarr/core/buffer/cpu.py

@@ -19,7 +19,7 @@
     from collections.abc import Callable, Iterable
     from typing import Self
 
-    from zarr.core.buffer.core import ArrayLike, NDArrayLike
+    from zarr.core.buffer.core import ArrayLike, NDArrayOrScalarLike
     from zarr.core.common import BytesLike
 
 
@@ -142,7 +142,7 @@ class NDBuffer(core.NDBuffer):
         ndarray-like object that is convertible to a regular Numpy array.
     """
 
-    def __init__(self, array: NDArrayLike) -> None:
+    def __init__(self, array: NDArrayOrScalarLike) -> None:
         super().__init__(array)
 
     @classmethod