diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index b42e036cfa..72f65a073b 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -59,7 +59,7 @@ jobs: conda activate zarr-env python -m pip install --upgrade pip python -m pip install -U pip setuptools wheel codecov line_profiler - python -m pip install -rrequirements_dev_minimal.txt numpy${{ matrix.numpy_version}} -rrequirements_dev_optional.txt pymongo redis + python -m pip install -rrequirements_dev_minimal.txt numpy${{matrix.numpy_version}} -rrequirements_dev_optional.txt pymongo redis python -m pip install . python -m pip freeze - name: Tests diff --git a/zarr/core.py b/zarr/core.py index bd61639ef6..e5b2045160 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -4,12 +4,11 @@ import math import operator import re -from collections.abc import MutableMapping from functools import reduce from typing import Any import numpy as np -from numcodecs.compat import ensure_bytes, ensure_ndarray +from numcodecs.compat import ensure_bytes from zarr._storage.store import _prefix_to_attrs_key, assert_zarr_v3_api_available from zarr.attrs import Attributes @@ -35,6 +34,7 @@ from zarr.storage import ( _get_hierarchy_metadata, _prefix_to_array_key, + KVStore, getsize, listdir, normalize_store_arg, @@ -51,6 +51,7 @@ normalize_shape, normalize_storage_path, PartialReadBuffer, + ensure_ndarray_like ) @@ -98,6 +99,12 @@ class Array: .. versionadded:: 2.11 + meta_array : array-like, optional + An array instance to use for determining arrays to create and return + to users. Use `numpy.empty(())` by default. + + .. versionadded:: 2.13 + Attributes ---------- @@ -129,6 +136,7 @@ class Array: vindex oindex write_empty_chunks + meta_array Methods ------- @@ -163,6 +171,7 @@ def __init__( partial_decompress=False, write_empty_chunks=True, zarr_version=None, + meta_array=None, ): # N.B., expect at this point store is fully initialized with all # configuration metadata fully specified and normalized @@ -191,8 +200,11 @@ def __init__( self._is_view = False self._partial_decompress = partial_decompress self._write_empty_chunks = write_empty_chunks + if meta_array is not None: + self._meta_array = np.empty_like(meta_array, shape=()) + else: + self._meta_array = np.empty(()) self._version = zarr_version - if self._version == 3: self._data_key_prefix = 'data/root/' + self._key_prefix self._data_path = 'data/root/' + self._path @@ -555,6 +567,13 @@ def write_empty_chunks(self) -> bool: """ return self._write_empty_chunks + @property + def meta_array(self): + """An array-like instance to use for determining arrays to create and return + to users. + """ + return self._meta_array + def __eq__(self, other): return ( isinstance(other, Array) and @@ -929,7 +948,7 @@ def _get_basic_selection_zd(self, selection, out=None, fields=None): except KeyError: # chunk not initialized - chunk = np.zeros((), dtype=self._dtype) + chunk = np.zeros_like(self._meta_array, shape=(), dtype=self._dtype) if self._fill_value is not None: chunk.fill(self._fill_value) @@ -1233,7 +1252,8 @@ def _get_selection(self, indexer, out=None, fields=None): # setup output array if out is None: - out = np.empty(out_shape, dtype=out_dtype, order=self._order) + out = np.empty_like(self._meta_array, shape=out_shape, + dtype=out_dtype, order=self._order) else: check_array_shape('out', out, out_shape) @@ -1607,9 +1627,13 @@ def set_coordinate_selection(self, selection, value, fields=None): # setup indexer indexer = CoordinateIndexer(selection, self) - # handle value - need to flatten + # handle value - need ndarray-like flatten value if not is_scalar(value, self._dtype): - value = np.asanyarray(value) + try: + value = ensure_ndarray_like(value) + except TypeError: + # Handle types like `list` or `tuple` + value = np.array(value, like=self._meta_array) if hasattr(value, 'shape') and len(value.shape) > 1: value = value.reshape(-1) @@ -1712,7 +1736,7 @@ def _set_basic_selection_zd(self, selection, value, fields=None): except KeyError: # chunk not initialized - chunk = np.zeros((), dtype=self._dtype) + chunk = np.zeros_like(self._meta_array, shape=(), dtype=self._dtype) if self._fill_value is not None: chunk.fill(self._fill_value) @@ -1772,7 +1796,7 @@ def _set_selection(self, indexer, value, fields=None): pass else: if not hasattr(value, 'shape'): - value = np.asanyarray(value) + value = np.asanyarray(value, like=self._meta_array) check_array_shape('value', value, sel_shape) # iterate over chunks in range @@ -1840,8 +1864,11 @@ def _process_chunk( self._dtype != object): dest = out[out_selection] + # Assume that array-like objects that doesn't have a + # `writeable` flag is writable. + dest_is_writable = getattr(dest, "writeable", True) write_direct = ( - dest.flags.writeable and + dest_is_writable and ( (self._order == 'C' and dest.flags.c_contiguous) or (self._order == 'F' and dest.flags.f_contiguous) @@ -1858,7 +1885,7 @@ def _process_chunk( cdata = cdata.read_full() self._compressor.decode(cdata, dest) else: - chunk = ensure_ndarray(cdata).view(self._dtype) + chunk = ensure_ndarray_like(cdata).view(self._dtype) chunk = chunk.reshape(self._chunks, order=self._order) np.copyto(dest, chunk) return @@ -1868,7 +1895,7 @@ def _process_chunk( if partial_read_decode: cdata.prepare_chunk() # size of chunk - tmp = np.empty(self._chunks, dtype=self.dtype) + tmp = np.empty_like(self._meta_array, shape=self._chunks, dtype=self.dtype) index_selection = PartialChunkIterator(chunk_selection, self.chunks) for start, nitems, partial_out_selection in index_selection: expected_shape = [ @@ -1925,7 +1952,7 @@ def _chunk_getitem(self, chunk_coords, chunk_selection, out, out_selection, """ out_is_ndarray = True try: - out = ensure_ndarray(out) + out = ensure_ndarray_like(out) except TypeError: out_is_ndarray = False @@ -1960,7 +1987,7 @@ def _chunk_getitems(self, lchunk_coords, lchunk_selection, out, lout_selection, """ out_is_ndarray = True try: - out = ensure_ndarray(out) + out = ensure_ndarray_like(out) except TypeError: # pragma: no cover out_is_ndarray = False @@ -2082,7 +2109,9 @@ def _process_for_setitem(self, ckey, chunk_selection, value, fields=None): if is_scalar(value, self._dtype): # setup array filled with value - chunk = np.empty(self._chunks, dtype=self._dtype, order=self._order) + chunk = np.empty_like( + self._meta_array, shape=self._chunks, dtype=self._dtype, order=self._order + ) chunk.fill(value) else: @@ -2102,14 +2131,18 @@ def _process_for_setitem(self, ckey, chunk_selection, value, fields=None): # chunk not initialized if self._fill_value is not None: - chunk = np.empty(self._chunks, dtype=self._dtype, order=self._order) + chunk = np.empty_like( + self._meta_array, shape=self._chunks, dtype=self._dtype, order=self._order + ) chunk.fill(self._fill_value) elif self._dtype == object: chunk = np.empty(self._chunks, dtype=self._dtype, order=self._order) else: # N.B., use zeros here so any region beyond the array has consistent # and compressible data - chunk = np.zeros(self._chunks, dtype=self._dtype, order=self._order) + chunk = np.zeros_like( + self._meta_array, shape=self._chunks, dtype=self._dtype, order=self._order + ) else: @@ -2159,7 +2192,7 @@ def _decode_chunk(self, cdata, start=None, nitems=None, expected_shape=None): chunk = f.decode(chunk) # view as numpy array with correct dtype - chunk = ensure_ndarray(chunk) + chunk = ensure_ndarray_like(chunk) # special case object dtype, because incorrect handling can lead to # segfaults and other bad things happening if self._dtype != object: @@ -2186,7 +2219,7 @@ def _encode_chunk(self, chunk): chunk = f.encode(chunk) # check object encoding - if ensure_ndarray(chunk).dtype == object: + if ensure_ndarray_like(chunk).dtype == object: raise RuntimeError('cannot write object array without object codec') # compress @@ -2196,7 +2229,7 @@ def _encode_chunk(self, chunk): cdata = chunk # ensure in-memory data is immutable and easy to compare - if isinstance(self.chunk_store, MutableMapping): + if isinstance(self.chunk_store, KVStore): cdata = ensure_bytes(cdata) return cdata @@ -2354,12 +2387,22 @@ def hexdigest(self, hashname="sha1"): return checksum def __getstate__(self): - return (self._store, self._path, self._read_only, self._chunk_store, - self._synchronizer, self._cache_metadata, self._attrs.cache, - self._partial_decompress, self._write_empty_chunks, self._version) + return { + "store": self._store, + "path": self._path, + "read_only": self._read_only, + "chunk_store": self._chunk_store, + "synchronizer": self._synchronizer, + "cache_metadata": self._cache_metadata, + "cache_attrs": self._attrs.cache, + "partial_decompress": self._partial_decompress, + "write_empty_chunks": self._write_empty_chunks, + "zarr_version": self._version, + "meta_array": self._meta_array, + } def __setstate__(self, state): - self.__init__(*state) + self.__init__(**state) def _synchronized_op(self, f, *args, **kwargs): @@ -2466,7 +2509,7 @@ def append(self, data, axis=0): Parameters ---------- - data : array_like + data : array-like Data to be appended. axis : int Axis along which to append. @@ -2502,7 +2545,7 @@ def _append_nosync(self, data, axis=0): # ensure data is array-like if not hasattr(data, 'shape'): - data = np.asanyarray(data) + data = np.asanyarray(data, like=self._meta_array) # ensure shapes are compatible for non-append dimensions self_shape_preserved = tuple(s for i, s in enumerate(self._shape) diff --git a/zarr/creation.py b/zarr/creation.py index e77f26b3e2..e1c815ed21 100644 --- a/zarr/creation.py +++ b/zarr/creation.py @@ -21,7 +21,7 @@ def create(shape, chunks=True, dtype=None, compressor='default', overwrite=False, path=None, chunk_store=None, filters=None, cache_metadata=True, cache_attrs=True, read_only=False, object_codec=None, dimension_separator=None, write_empty_chunks=True, - *, zarr_version=None, **kwargs): + *, zarr_version=None, meta_array=None, **kwargs): """Create an array. Parameters @@ -89,6 +89,14 @@ def create(shape, chunks=True, dtype=None, compressor='default', inferred from ``store`` or ``chunk_store`` if they are provided, otherwise defaulting to 2. + .. versionadded:: 2.12 + + meta_array : array-like, optional + An array instance to use for determining arrays to create and return + to users. Use `numpy.empty(())` by default. + + .. versionadded:: 2.13 + Returns ------- z : zarr.core.Array @@ -166,7 +174,7 @@ def create(shape, chunks=True, dtype=None, compressor='default', # instantiate array z = Array(store, path=path, chunk_store=chunk_store, synchronizer=synchronizer, cache_metadata=cache_metadata, cache_attrs=cache_attrs, read_only=read_only, - write_empty_chunks=write_empty_chunks) + write_empty_chunks=write_empty_chunks, meta_array=meta_array) return z diff --git a/zarr/hierarchy.py b/zarr/hierarchy.py index 80da3ddbc6..177d1eec71 100644 --- a/zarr/hierarchy.py +++ b/zarr/hierarchy.py @@ -64,6 +64,12 @@ class Group(MutableMapping): synchronizer : object, optional Array synchronizer. + meta_array : array-like, optional + An array instance to use for determining arrays to create and return + to users. Use `numpy.empty(())` by default. + + .. versionadded:: 2.13 + Attributes ---------- store @@ -74,6 +80,7 @@ class Group(MutableMapping): synchronizer attrs info + meta_array Methods ------- @@ -114,7 +121,8 @@ class Group(MutableMapping): """ def __init__(self, store, path=None, read_only=False, chunk_store=None, - cache_attrs=True, synchronizer=None, zarr_version=None): + cache_attrs=True, synchronizer=None, zarr_version=None, *, + meta_array=None): store: BaseStore = _normalize_store_arg(store, zarr_version=zarr_version) if zarr_version is None: zarr_version = getattr(store, '_store_version', DEFAULT_ZARR_VERSION) @@ -133,8 +141,11 @@ def __init__(self, store, path=None, read_only=False, chunk_store=None, self._key_prefix = '' self._read_only = read_only self._synchronizer = synchronizer + if meta_array is not None: + self._meta_array = np.empty_like(meta_array, shape=()) + else: + self._meta_array = np.empty(()) self._version = zarr_version - if self._version == 3: self._data_key_prefix = data_root + self._key_prefix self._data_path = data_root + self._path @@ -231,6 +242,13 @@ def info(self): """Return diagnostic information about the group.""" return self._info + @property + def meta_array(self): + """An array-like instance to use for determining arrays to create and return + to users. + """ + return self._meta_array + def __eq__(self, other): return ( isinstance(other, Group) and @@ -351,11 +369,19 @@ def typestr(o): return items def __getstate__(self): - return (self._store, self._path, self._read_only, self._chunk_store, - self.attrs.cache, self._synchronizer) + return { + "store": self._store, + "path": self._path, + "read_only": self._read_only, + "chunk_store": self._chunk_store, + "cache_attrs": self._attrs.cache, + "synchronizer": self._synchronizer, + "zarr_version": self._version, + "meta_array": self._meta_array, + } def __setstate__(self, state): - self.__init__(*state) + self.__init__(**state) def _item_path(self, item): absolute = isinstance(item, str) and item and item[0] == '/' @@ -411,18 +437,20 @@ def __getitem__(self, item): return Array(self._store, read_only=self._read_only, path=path, chunk_store=self._chunk_store, synchronizer=self._synchronizer, cache_attrs=self.attrs.cache, - zarr_version=self._version) + zarr_version=self._version, meta_array=self._meta_array) elif contains_group(self._store, path, explicit_only=True): return Group(self._store, read_only=self._read_only, path=path, chunk_store=self._chunk_store, cache_attrs=self.attrs.cache, - synchronizer=self._synchronizer, zarr_version=self._version) + synchronizer=self._synchronizer, zarr_version=self._version, + meta_array=self._meta_array) elif self._version == 3: implicit_group = meta_root + path + '/' # non-empty folder in the metadata path implies an implicit group if self._store.list_prefix(implicit_group): return Group(self._store, read_only=self._read_only, path=path, chunk_store=self._chunk_store, cache_attrs=self.attrs.cache, - synchronizer=self._synchronizer, zarr_version=self._version) + synchronizer=self._synchronizer, zarr_version=self._version, + meta_array=self._meta_array) else: raise KeyError(item) else: @@ -895,7 +923,7 @@ def create_dataset(self, name, **kwargs): ---------- name : string Array name. - data : array_like, optional + data : array-like, optional Initial data. shape : int or tuple of ints Array shape. @@ -1006,7 +1034,8 @@ def _require_dataset_nosync(self, name, shape, dtype=None, exact=False, cache_attrs = kwargs.get('cache_attrs', self.attrs.cache) a = Array(self._store, path=path, read_only=self._read_only, chunk_store=self._chunk_store, synchronizer=synchronizer, - cache_metadata=cache_metadata, cache_attrs=cache_attrs) + cache_metadata=cache_metadata, cache_attrs=cache_attrs, + meta_array=self._meta_array) shape = normalize_shape(shape) if shape != a.shape: raise TypeError('shape do not match existing array; expected {}, got {}' @@ -1266,7 +1295,7 @@ def group(store=None, overwrite=False, chunk_store=None, def open_group(store=None, mode='a', cache_attrs=True, synchronizer=None, path=None, - chunk_store=None, storage_options=None, *, zarr_version=None): + chunk_store=None, storage_options=None, *, zarr_version=None, meta_array=None): """Open a group using file-mode-like semantics. Parameters @@ -1291,6 +1320,11 @@ def open_group(store=None, mode='a', cache_attrs=True, synchronizer=None, path=N storage_options : dict If using an fsspec URL to create the store, these will be passed to the backend implementation. Ignored otherwise. + meta_array : array-like, optional + An array instance to use for determining arrays to create and return + to users. Use `numpy.empty(())` by default. + + .. versionadded:: 2.13 Returns ------- @@ -1368,4 +1402,4 @@ def open_group(store=None, mode='a', cache_attrs=True, synchronizer=None, path=N return Group(store, read_only=read_only, cache_attrs=cache_attrs, synchronizer=synchronizer, path=path, chunk_store=chunk_store, - zarr_version=zarr_version) + zarr_version=zarr_version, meta_array=meta_array) diff --git a/zarr/storage.py b/zarr/storage.py index eb5106078b..4a1408ec01 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -39,7 +39,7 @@ from numcodecs.compat import ( ensure_bytes, ensure_text, - ensure_contiguous_ndarray + ensure_contiguous_ndarray_like ) from numcodecs.registry import codec_registry @@ -55,7 +55,8 @@ from zarr.util import (buffer_size, json_loads, nolock, normalize_chunks, normalize_dimension_separator, normalize_dtype, normalize_fill_value, normalize_order, - normalize_shape, normalize_storage_path, retry_call) + normalize_shape, normalize_storage_path, retry_call + ) from zarr._storage.absstore import ABSStore # noqa: F401 from zarr._storage.store import (_get_hierarchy_metadata, # noqa: F401 @@ -1070,7 +1071,7 @@ def __setitem__(self, key, value): key = self._normalize_key(key) # coerce to flat, contiguous array (ideally without copying) - value = ensure_contiguous_ndarray(value) + value = ensure_contiguous_ndarray_like(value) # destination path for key file_path = os.path.join(self.path, key) @@ -1755,7 +1756,7 @@ def __getitem__(self, key): def __setitem__(self, key, value): if self.mode == 'r': raise ReadOnlyError() - value = ensure_contiguous_ndarray(value).view("u1") + value = ensure_contiguous_ndarray_like(value).view("u1") with self.mutex: # writestr(key, value) writes with default permissions from # zipfile (600) that are too restrictive, build ZipInfo for @@ -2601,7 +2602,7 @@ def update(self, *args, **kwargs): kv_list = [] for dct in args: for k, v in dct.items(): - v = ensure_contiguous_ndarray(v) + v = ensure_contiguous_ndarray_like(v) # Accumulate key-value pairs for storage kv_list.append((k, v)) diff --git a/zarr/tests/test_meta_array.py b/zarr/tests/test_meta_array.py new file mode 100644 index 0000000000..6172af3be9 --- /dev/null +++ b/zarr/tests/test_meta_array.py @@ -0,0 +1,233 @@ +from typing import Optional +import numpy as np +import pytest + +from numcodecs.abc import Codec +from numcodecs.compat import ensure_contiguous_ndarray_like +from numcodecs.registry import get_codec, register_codec + +import zarr.codecs +from zarr.core import Array +from zarr.creation import array, empty, full, ones, zeros +from zarr.hierarchy import open_group +from zarr.storage import DirectoryStore, MemoryStore, Store, ZipStore + + +class CuPyCPUCompressor(Codec): # pragma: no cover + """CPU compressor for CuPy arrays + + This compressor converts CuPy arrays host memory before compressing + the arrays using `compressor`. + + Parameters + ---------- + compressor : numcodecs.abc.Codec + The codec to use for compression and decompression. + """ + + codec_id = "cupy_cpu_compressor" + + def __init__(self, compressor: Codec = None): + self.compressor = compressor + + def encode(self, buf): + import cupy + + buf = cupy.asnumpy(ensure_contiguous_ndarray_like(buf)) + if self.compressor: + buf = self.compressor.encode(buf) + return buf + + def decode(self, chunk, out=None): + import cupy + + if self.compressor: + cpu_out = None if out is None else cupy.asnumpy(out) + chunk = self.compressor.decode(chunk, cpu_out) + + chunk = cupy.asarray(ensure_contiguous_ndarray_like(chunk)) + if out is not None: + cupy.copyto(out, chunk.view(dtype=out.dtype), casting="no") + chunk = out + return chunk + + def get_config(self): + cc_config = self.compressor.get_config() if self.compressor else None + return { + "id": self.codec_id, + "compressor_config": cc_config, + } + + @classmethod + def from_config(cls, config): + cc_config = config.get("compressor_config", None) + compressor = get_codec(cc_config) if cc_config else None + return cls(compressor=compressor) + + +register_codec(CuPyCPUCompressor) + + +class MyArray(np.ndarray): + """Dummy array class to test the `meta_array` argument + + Useful when CuPy isn't available. + + This class also makes some of the functions from the numpy + module available. + """ + + testing = np.testing + + @classmethod + def arange(cls, size): + ret = cls(shape=(size,), dtype="int64") + ret[:] = range(size) + return ret + + @classmethod + def empty(cls, shape): + return cls(shape=shape) + + +def init_compressor(compressor) -> CuPyCPUCompressor: + if compressor: + compressor = getattr(zarr.codecs, compressor)() + return CuPyCPUCompressor(compressor) + + +def init_store(tmp_path, store_type) -> Optional[Store]: + if store_type is DirectoryStore: + return store_type(str(tmp_path / "store")) + if store_type is MemoryStore: + return MemoryStore() + return None + + +def ensure_module(module): + if isinstance(module, str): + return pytest.importorskip(module) + return module + + +param_module_and_compressor = [ + (MyArray, None), + ("cupy", init_compressor(None)), + ("cupy", init_compressor("Zlib")), + ("cupy", init_compressor("Blosc")), +] + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +@pytest.mark.parametrize("store_type", [None, DirectoryStore, MemoryStore, ZipStore]) +def test_array(tmp_path, module, compressor, store_type): + xp = ensure_module(module) + + store = init_store(tmp_path / "from_cupy_array", store_type) + a = xp.arange(100) + z = array(a, chunks=10, compressor=compressor, store=store, meta_array=xp.empty(())) + assert a.shape == z.shape + assert a.dtype == z.dtype + assert isinstance(a, type(z[:])) + assert isinstance(z.meta_array, type(xp.empty(()))) + xp.testing.assert_array_equal(a, z[:]) + + # with array-like + store = init_store(tmp_path / "from_list", store_type) + a = list(range(100)) + z = array(a, chunks=10, compressor=compressor, store=store, meta_array=xp.empty(())) + assert (100,) == z.shape + assert np.asarray(a).dtype == z.dtype + xp.testing.assert_array_equal(a, z[:]) + + # with another zarr array + store = init_store(tmp_path / "from_another_store", store_type) + z2 = array(z, compressor=compressor, store=store, meta_array=xp.empty(())) + assert z.shape == z2.shape + assert z.chunks == z2.chunks + assert z.dtype == z2.dtype + xp.testing.assert_array_equal(z[:], z2[:]) + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +def test_empty(module, compressor): + xp = ensure_module(module) + z = empty( + 100, + chunks=10, + compressor=compressor, + meta_array=xp.empty(()), + ) + assert (100,) == z.shape + assert (10,) == z.chunks + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +def test_zeros(module, compressor): + xp = ensure_module(module) + z = zeros( + 100, + chunks=10, + compressor=compressor, + meta_array=xp.empty(()), + ) + assert (100,) == z.shape + assert (10,) == z.chunks + xp.testing.assert_array_equal(np.zeros(100), z[:]) + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +def test_ones(module, compressor): + xp = ensure_module(module) + z = ones( + 100, + chunks=10, + compressor=compressor, + meta_array=xp.empty(()), + ) + assert (100,) == z.shape + assert (10,) == z.chunks + xp.testing.assert_array_equal(np.ones(100), z[:]) + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +def test_full(module, compressor): + xp = ensure_module(module) + z = full( + 100, + chunks=10, + fill_value=42, + dtype="i4", + compressor=compressor, + meta_array=xp.empty(()), + ) + assert (100,) == z.shape + assert (10,) == z.chunks + xp.testing.assert_array_equal(np.full(100, fill_value=42, dtype="i4"), z[:]) + + # nan + z = full( + 100, + chunks=10, + fill_value=np.nan, + dtype="f8", + compressor=compressor, + meta_array=xp.empty(()), + ) + assert np.all(np.isnan(z[:])) + + +@pytest.mark.parametrize("module, compressor", param_module_and_compressor) +@pytest.mark.parametrize("store_type", [None, DirectoryStore, MemoryStore, ZipStore]) +def test_group(tmp_path, module, compressor, store_type): + xp = ensure_module(module) + store = init_store(tmp_path, store_type) + g = open_group(store, meta_array=xp.empty(())) + g.ones("data", shape=(10, 11), dtype=int, compressor=compressor) + a = g["data"] + assert a.shape == (10, 11) + assert a.dtype == int + assert isinstance(a, Array) + assert isinstance(a[:], type(xp.empty(()))) + assert (a[:] == 1).all() + assert isinstance(g.meta_array, type(xp.empty(()))) diff --git a/zarr/util.py b/zarr/util.py index cc3bd50356..c9136a63eb 100644 --- a/zarr/util.py +++ b/zarr/util.py @@ -10,7 +10,7 @@ from asciitree import BoxStyle, LeftAligned from asciitree.traversal import Traversal from collections.abc import Iterable -from numcodecs.compat import ensure_ndarray, ensure_text +from numcodecs.compat import ensure_text, ensure_ndarray_like from numcodecs.registry import codec_registry from numcodecs.blosc import cbuffer_sizes, cbuffer_metainfo @@ -352,7 +352,7 @@ def normalize_storage_path(path: Union[str, bytes, None]) -> str: def buffer_size(v) -> int: - return ensure_ndarray(v).nbytes + return ensure_ndarray_like(v).nbytes def info_text_report(items: Dict[Any, Any]) -> str: