Skip to content

Commit 888304a

Browse files
brokkoli71brokkoli71
committed
create NDArrayOrScalarLike
1 parent 862481a commit 888304a

File tree

10 files changed

+66
-51
lines changed

10 files changed

+66
-51
lines changed

‎src/zarr/api/asynchronous.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111

1212
from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
1313
from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
14-
from zarr.core.buffer import NDArrayLike
14+
from zarr.core.buffer import NDArrayOrScalarLike
1515
from zarr.core.common import (
1616
JSON,
1717
AccessModeLiteral,
@@ -232,7 +232,7 @@ async def load(
232232
path: str | None = None,
233233
zarr_format: ZarrFormat | None = None,
234234
zarr_version: ZarrFormat | None = None,
235-
) -> NDArrayLike | dict[str, NDArrayLike]:
235+
) -> NDArrayOrScalarLike | dict[str, NDArrayOrScalarLike]:
236236
"""Load data from an array or group into memory.
237237
238238
Parameters
@@ -348,7 +348,7 @@ async def open_consolidated(
348348

349349
async def save(
350350
store: StoreLike,
351-
*args: NDArrayLike,
351+
*args: NDArrayOrScalarLike,
352352
zarr_version: ZarrFormat | None = None, # deprecated
353353
zarr_format: ZarrFormat | None = None,
354354
path: str | None = None,
@@ -381,7 +381,7 @@ async def save(
381381

382382
async def save_array(
383383
store: StoreLike,
384-
arr: NDArrayLike,
384+
arr: NDArrayOrScalarLike,
385385
*,
386386
zarr_version: ZarrFormat | None = None, # deprecated
387387
zarr_format: ZarrFormat | None = None,
@@ -412,8 +412,8 @@ async def save_array(
412412
_handle_zarr_version_or_format(zarr_version=zarr_version, zarr_format=zarr_format)
413413
or _default_zarr_format()
414414
)
415-
if not isinstance(arr, NDArrayLike):
416-
raise TypeError("arr argument must be numpy or other NDArrayLike array")
415+
if not isinstance(arr, NDArrayOrScalarLike):
416+
raise TypeError("arr argument must be numpy or other NDArrayOrScalarLike array")
417417

418418
mode = kwargs.pop("mode", "a")
419419
store_path = await make_store_path(store, path=path, mode=mode, storage_options=storage_options)
@@ -436,12 +436,12 @@ async def save_array(
436436

437437
async def save_group(
438438
store: StoreLike,
439-
*args: NDArrayLike,
439+
*args: NDArrayOrScalarLike,
440440
zarr_version: ZarrFormat | None = None, # deprecated
441441
zarr_format: ZarrFormat | None = None,
442442
path: str | None = None,
443443
storage_options: dict[str, Any] | None = None,
444-
**kwargs: NDArrayLike,
444+
**kwargs: NDArrayOrScalarLike,
445445
) -> None:
446446
"""Convenience function to save several NumPy arrays to the local file system, following a
447447
similar API to the NumPy savez()/savez_compressed() functions.
@@ -474,13 +474,15 @@ async def save_group(
474474
)
475475

476476
for arg in args:
477-
if not isinstance(arg, NDArrayLike):
477+
if not isinstance(arg, NDArrayOrScalarLike):
478478
raise TypeError(
479-
"All arguments must be numpy or other NDArrayLike arrays (except store, path, storage_options, and zarr_format)"
479+
"All arguments must be numpy or other NDArrayOrScalarLike arrays (except store, path, storage_options, and zarr_format)"
480480
)
481481
for k, v in kwargs.items():
482-
if not isinstance(v, NDArrayLike):
483-
raise TypeError(f"Keyword argument '{k}' must be a numpy or other NDArrayLike array")
482+
if not isinstance(v, NDArrayOrScalarLike):
483+
raise TypeError(
484+
f"Keyword argument '{k}' must be a numpy or other NDArrayOrScalarLike array"
485+
)
484486

485487
if len(args) == 0 and len(kwargs) == 0:
486488
raise ValueError("at least one array must be provided")

‎src/zarr/api/synchronous.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
ShardsLike,
2727
)
2828
from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
29-
from zarr.core.buffer import NDArrayLike
29+
from zarr.core.buffer import NDArrayOrScalarLike
3030
from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
3131
from zarr.core.common import (
3232
JSON,
@@ -119,7 +119,7 @@ def load(
119119
path: str | None = None,
120120
zarr_format: ZarrFormat | None = None,
121121
zarr_version: ZarrFormat | None = None,
122-
) -> NDArrayLike | dict[str, NDArrayLike]:
122+
) -> NDArrayOrScalarLike | dict[str, NDArrayOrScalarLike]:
123123
"""Load data from an array or group into memory.
124124
125125
Parameters
@@ -216,7 +216,7 @@ def open_consolidated(*args: Any, use_consolidated: Literal[True] = True, **kwar
216216

217217
def save(
218218
store: StoreLike,
219-
*args: NDArrayLike,
219+
*args: NDArrayOrScalarLike,
220220
zarr_version: ZarrFormat | None = None, # deprecated
221221
zarr_format: ZarrFormat | None = None,
222222
path: str | None = None,
@@ -247,7 +247,7 @@ def save(
247247
@_deprecate_positional_args
248248
def save_array(
249249
store: StoreLike,
250-
arr: NDArrayLike,
250+
arr: NDArrayOrScalarLike,
251251
*,
252252
zarr_version: ZarrFormat | None = None, # deprecated
253253
zarr_format: ZarrFormat | None = None,
@@ -290,12 +290,12 @@ def save_array(
290290

291291
def save_group(
292292
store: StoreLike,
293-
*args: NDArrayLike,
293+
*args: NDArrayOrScalarLike,
294294
zarr_version: ZarrFormat | None = None, # deprecated
295295
zarr_format: ZarrFormat | None = None,
296296
path: str | None = None,
297297
storage_options: dict[str, Any] | None = None,
298-
**kwargs: NDArrayLike,
298+
**kwargs: NDArrayOrScalarLike,
299299
) -> None:
300300
"""Save several NumPy arrays to the local file system.
301301

‎src/zarr/codecs/bytes.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import numpy as np
99

1010
from zarr.abc.codec import ArrayBytesCodec
11-
from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
11+
from zarr.core.buffer import Buffer, NDArrayOrScalarLike, NDBuffer
1212
from zarr.core.common import JSON, parse_enum, parse_named_configuration
1313
from zarr.registry import register_codec
1414

@@ -81,7 +81,7 @@ async def _decode_single(
8181
dtype = np.dtype(f"|{chunk_spec.dtype.str[1:]}")
8282

8383
as_array_like = chunk_bytes.as_array_like()
84-
if isinstance(as_array_like, NDArrayLike):
84+
if isinstance(as_array_like, NDArrayOrScalarLike):
8585
as_nd_array_like = as_array_like
8686
else:
8787
as_nd_array_like = np.asanyarray(as_array_like)

‎src/zarr/core/array.py‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
from zarr.core.attributes import Attributes
3535
from zarr.core.buffer import (
3636
BufferPrototype,
37-
NDArrayLike,
37+
NDArrayOrScalarLike,
3838
NDBuffer,
3939
default_buffer_prototype,
4040
)
@@ -1255,7 +1255,7 @@ async def _get_selection(
12551255
prototype: BufferPrototype,
12561256
out: NDBuffer | None = None,
12571257
fields: Fields | None = None,
1258-
) -> NDArrayLike:
1258+
) -> NDArrayOrScalarLike:
12591259
# check fields are sensible
12601260
out_dtype = check_fields(fields, self.dtype)
12611261

@@ -1305,7 +1305,7 @@ async def getitem(
13051305
selection: BasicSelection,
13061306
*,
13071307
prototype: BufferPrototype | None = None,
1308-
) -> NDArrayLike:
1308+
) -> NDArrayOrScalarLike:
13091309
"""
13101310
Asynchronous function that retrieves a subset of the array's data based on the provided selection.
13111311
@@ -1318,7 +1318,7 @@ async def getitem(
13181318
13191319
Returns
13201320
-------
1321-
NDArrayLike
1321+
NDArrayOrScalarLike
13221322
The retrieved subset of the array's data.
13231323
13241324
Examples
@@ -1396,11 +1396,11 @@ async def _set_selection(
13961396
# ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}"
13971397
if not hasattr(value, "dtype") or value.dtype.name != self.metadata.dtype.name:
13981398
if hasattr(value, "astype"):
1399-
# Handle things that are already NDArrayLike more efficiently
1399+
# Handle things that are already NDArrayOrScalarLike more efficiently
14001400
value = value.astype(dtype=self.metadata.dtype, order="A")
14011401
else:
14021402
value = np.array(value, dtype=self.metadata.dtype, order="A")
1403-
value = cast(NDArrayLike, value)
1403+
value = cast(NDArrayOrScalarLike, value)
14041404
# We accept any ndarray like object from the user and convert it
14051405
# to a NDBuffer (or subclass). From this point onwards, we only pass
14061406
# Buffer and NDBuffer between components.
@@ -2260,7 +2260,7 @@ def _iter_chunk_regions(
22602260

22612261
def __array__(
22622262
self, dtype: npt.DTypeLike | None = None, copy: bool | None = None
2263-
) -> NDArrayLike:
2263+
) -> NDArrayOrScalarLike:
22642264
"""
22652265
This method is used by numpy when converting zarr.Array into a numpy array.
22662266
For more information, see https://numpy.org/devdocs/user/basics.interoperability.html#the-array-method
@@ -2278,7 +2278,7 @@ def __array__(
22782278

22792279
return arr_np
22802280

2281-
def __getitem__(self, selection: Selection) -> NDArrayLike:
2281+
def __getitem__(self, selection: Selection) -> NDArrayOrScalarLike:
22822282
"""Retrieve data for an item or region of the array.
22832283
22842284
Parameters
@@ -2289,7 +2289,7 @@ def __getitem__(self, selection: Selection) -> NDArrayLike:
22892289
22902290
Returns
22912291
-------
2292-
NDArrayLike
2292+
NDArrayOrScalarLike
22932293
An array-like containing the data for the requested region.
22942294
22952295
Examples
@@ -2536,7 +2536,7 @@ def get_basic_selection(
25362536
out: NDBuffer | None = None,
25372537
prototype: BufferPrototype | None = None,
25382538
fields: Fields | None = None,
2539-
) -> NDArrayLike:
2539+
) -> NDArrayOrScalarLike:
25402540
"""Retrieve data for an item or region of the array.
25412541
25422542
Parameters
@@ -2554,7 +2554,7 @@ def get_basic_selection(
25542554
25552555
Returns
25562556
-------
2557-
NDArrayLike
2557+
NDArrayOrScalarLike
25582558
An array-like containing the data for the requested region.
25592559
25602560
Examples
@@ -2756,7 +2756,7 @@ def get_orthogonal_selection(
27562756
out: NDBuffer | None = None,
27572757
fields: Fields | None = None,
27582758
prototype: BufferPrototype | None = None,
2759-
) -> NDArrayLike:
2759+
) -> NDArrayOrScalarLike:
27602760
"""Retrieve data by making a selection for each dimension of the array. For
27612761
example, if an array has 2 dimensions, allows selecting specific rows and/or
27622762
columns. The selection for each dimension can be either an integer (indexing a
@@ -2778,7 +2778,7 @@ def get_orthogonal_selection(
27782778
27792779
Returns
27802780
-------
2781-
NDArrayLike
2781+
NDArrayOrScalarLike
27822782
An array-like containing the data for the requested selection.
27832783
27842784
Examples
@@ -2992,7 +2992,7 @@ def get_mask_selection(
29922992
out: NDBuffer | None = None,
29932993
fields: Fields | None = None,
29942994
prototype: BufferPrototype | None = None,
2995-
) -> NDArrayLike:
2995+
) -> NDArrayOrScalarLike:
29962996
"""Retrieve a selection of individual items, by providing a Boolean array of the
29972997
same shape as the array against which the selection is being made, where True
29982998
values indicate a selected item.
@@ -3012,7 +3012,7 @@ def get_mask_selection(
30123012
30133013
Returns
30143014
-------
3015-
NDArrayLike
3015+
NDArrayOrScalarLike
30163016
An array-like containing the data for the requested selection.
30173017
30183018
Examples
@@ -3154,7 +3154,7 @@ def get_coordinate_selection(
31543154
out: NDBuffer | None = None,
31553155
fields: Fields | None = None,
31563156
prototype: BufferPrototype | None = None,
3157-
) -> NDArrayLike:
3157+
) -> NDArrayOrScalarLike:
31583158
"""Retrieve a selection of individual items, by providing the indices
31593159
(coordinates) for each selected item.
31603160
@@ -3172,7 +3172,7 @@ def get_coordinate_selection(
31723172
31733173
Returns
31743174
-------
3175-
NDArrayLike
3175+
NDArrayOrScalarLike
31763176
An array-like containing the data for the requested coordinate selection.
31773177
31783178
Examples
@@ -3325,7 +3325,7 @@ def set_coordinate_selection(
33253325
value = np.array(value).reshape(-1)
33263326

33273327
if not is_scalar(value, self.dtype) and (
3328-
isinstance(value, NDArrayLike) and indexer.shape != value.shape
3328+
isinstance(value, NDArrayOrScalarLike) and indexer.shape != value.shape
33293329
):
33303330
raise ValueError(
33313331
f"Attempting to set a selection of {indexer.sel_shape[0]} "
@@ -3342,7 +3342,7 @@ def get_block_selection(
33423342
out: NDBuffer | None = None,
33433343
fields: Fields | None = None,
33443344
prototype: BufferPrototype | None = None,
3345-
) -> NDArrayLike:
3345+
) -> NDArrayOrScalarLike:
33463346
"""Retrieve a selection of individual items, by providing the indices
33473347
(coordinates) for each selected item.
33483348
@@ -3360,7 +3360,7 @@ def get_block_selection(
33603360
33613361
Returns
33623362
-------
3363-
NDArrayLike
3363+
NDArrayOrScalarLike
33643364
An array-like containing the data for the requested block selection.
33653365
33663366
Examples

‎src/zarr/core/buffer/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
ArrayLike,
33
Buffer,
44
BufferPrototype,
5-
NDArrayLike,
5+
NDArrayOrScalarLike,
66
NDBuffer,
77
default_buffer_prototype,
88
)
@@ -12,7 +12,7 @@
1212
"ArrayLike",
1313
"Buffer",
1414
"BufferPrototype",
15-
"NDArrayLike",
15+
"NDArrayOrScalarLike",
1616
"NDBuffer",
1717
"default_buffer_prototype",
1818
"numpy_buffer_prototype",

‎src/zarr/core/buffer/core.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
NamedTuple,
1010
Protocol,
1111
SupportsIndex,
12+
Union,
1213
cast,
1314
runtime_checkable,
1415
)
@@ -106,6 +107,9 @@ def __eq__(self, other: object) -> Self: # type: ignore[explicit-override, over
106107
"""
107108

108109

110+
NDArrayLike = Union["np.ScalarType", NDArrayLike]
111+
112+
109113
class ScalarWrapper:
110114
def __init__(self, value: Any, dtype: np.dtype[Any] | None = None) -> None:
111115
self._value: Any = value
@@ -552,7 +556,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
552556
"""
553557
...
554558

555-
def as_scalar(self) -> np.generic:
559+
def as_scalar(self) -> np.ScalarType:
556560
"""Returns the buffer as a scalar value
557561
558562
Returns

‎src/zarr/core/buffer/cpu.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from collections.abc import Callable, Iterable
2020
from typing import Self
2121

22-
from zarr.core.buffer.core import ArrayLike, NDArrayLike
22+
from zarr.core.buffer.core import ArrayLike, NDArrayOrScalarLike
2323
from zarr.core.common import BytesLike
2424

2525

@@ -142,7 +142,7 @@ class NDBuffer(core.NDBuffer):
142142
ndarray-like object that is convertible to a regular Numpy array.
143143
"""
144144

145-
def __init__(self, array: NDArrayLike) -> None:
145+
def __init__(self, array: NDArrayOrScalarLike) -> None:
146146
super().__init__(array)
147147

148148
@classmethod

0 commit comments

Comments
 (0)