From bf57a711e1b87eff1691c6ce7f4f0b6eca968d42 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 5 Aug 2024 20:21:33 +0200 Subject: [PATCH 1/4] add type hints to indexing tests --- tests/v3/test_indexing.py | 155 +++++++++++++++++++++++--------------- 1 file changed, 95 insertions(+), 60 deletions(-) diff --git a/tests/v3/test_indexing.py b/tests/v3/test_indexing.py index 12547765a9..c84cc3eb24 100644 --- a/tests/v3/test_indexing.py +++ b/tests/v3/test_indexing.py @@ -9,12 +9,14 @@ import numpy.typing as npt import pytest from numpy.testing import assert_array_equal +from typing_extensions import Self import zarr -from zarr.abc.store import Store -from zarr.buffer import BufferPrototype, NDBuffer +from zarr.array import Array +from zarr.buffer import Buffer, BufferPrototype, NDBuffer from zarr.common import ChunkCoords from zarr.indexing import ( + BasicSelection, make_slice_selection, normalize_integer_selection, oindex, @@ -27,8 +29,8 @@ @pytest.fixture -async def store() -> Iterator[Store]: - yield StorePath(await MemoryStore.open(mode="w")) +def store() -> Iterator[StorePath]: + yield StorePath(MemoryStore(mode="w")) def zarr_array_from_numpy_array( @@ -48,24 +50,31 @@ def zarr_array_from_numpy_array( class CountingDict(MemoryStore): + counter: Counter[Any] + @classmethod - async def open(cls): + async def open(cls) -> Self: store = await super().open(mode="w") store.counter = Counter() return store - async def get(self, key, prototype: BufferPrototype, byte_range=None): + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: tuple[int | None, int | None] | None = None, + ) -> Buffer | None: key_suffix = "/".join(key.split("/")[1:]) self.counter["__getitem__", key_suffix] += 1 return await super().get(key, prototype, byte_range) - async def set(self, key, value, byte_range=None): + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: key_suffix = "/".join(key.split("/")[1:]) self.counter["__setitem__", key_suffix] += 1 return await super().set(key, value, byte_range) -def test_normalize_integer_selection(): +def test_normalize_integer_selection() -> None: assert 1 == normalize_integer_selection(1, 100) assert 99 == normalize_integer_selection(-1, 100) with pytest.raises(IndexError): @@ -76,7 +85,7 @@ def test_normalize_integer_selection(): normalize_integer_selection(-1000, 100) -def test_replace_ellipsis(): +def test_replace_ellipsis() -> None: # 1D, single item assert (0,) == replace_ellipsis(0, (100,)) @@ -165,7 +174,7 @@ def test_get_basic_selection_0d(store: StorePath, use_out: bool, value: Any, dty # assert_array_equal(a[["foo", "bar"]], c) -basic_selections_1d = [ +basic_selections_1d: list[BasicSelection] = [ # single value 42, -1, @@ -207,7 +216,7 @@ def test_get_basic_selection_0d(store: StorePath, use_out: bool, value: Any, dty slice(50, 150, 10), ] -basic_selections_1d_bad = [ +basic_selections_1d_bad: list[Any] = [ # only positive step supported slice(None, None, -1), slice(None, None, -10), @@ -239,7 +248,7 @@ def test_get_basic_selection_0d(store: StorePath, use_out: bool, value: Any, dty ] -def _test_get_basic_selection(a, z, selection): +def _test_get_basic_selection(a: Array, z: Array, selection: BasicSelection) -> None: expect = a[selection] actual = z.get_basic_selection(selection) assert_array_equal(expect, actual) @@ -253,7 +262,7 @@ def _test_get_basic_selection(a, z, selection): # noinspection PyStatementEffect -def test_get_basic_selection_1d(store: StorePath): +def test_get_basic_selection_1d(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -323,7 +332,7 @@ def test_get_basic_selection_1d(store: StorePath): # noinspection PyStatementEffect -def test_get_basic_selection_2d(store: StorePath): +def test_get_basic_selection_2d(store: StorePath) -> None: # setup a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) @@ -344,7 +353,7 @@ def test_get_basic_selection_2d(store: StorePath): np.testing.assert_array_equal(z[fancy_selection], [0, 11]) -def test_fancy_indexing_fallback_on_get_setitem(store: StorePath): +def test_fancy_indexing_fallback_on_get_setitem(store: StorePath) -> None: z = zarr_array_from_numpy_array(store, np.zeros((20, 20))) z[[1, 2, 3], [1, 2, 3]] = 1 np.testing.assert_array_equal( @@ -384,7 +393,9 @@ def test_fancy_indexing_fallback_on_get_setitem(store: StorePath): (([1, 0, 1]), [[3, 4, 5], [0, 1, 2], [3, 4, 5]]), ], ) -def test_orthogonal_indexing_fallback_on_getitem_2d(store: StorePath, index, expected_result): +def test_orthogonal_indexing_fallback_on_getitem_2d( + store: StorePath, index: BasicSelection, expected_result: list[list[int]] +) -> None: """ Tests the orthogonal indexing fallback on __getitem__ for a 2D matrix. @@ -414,7 +425,9 @@ def test_orthogonal_indexing_fallback_on_getitem_2d(store: StorePath, index, exp ((slice(0, 2), [1, 2], slice(0, 2)), [[[3, 4], [6, 7]], [[12, 13], [15, 16]]]), ], ) -def test_orthogonal_indexing_fallback_on_getitem_3d(store: StorePath, index, expected_result): +def test_orthogonal_indexing_fallback_on_getitem_3d( + store: StorePath, index: BasicSelection, expected_result: list[list[int]] +) -> None: """ Tests the orthogonal indexing fallback on __getitem__ for a 3D matrix. @@ -452,7 +465,9 @@ def test_orthogonal_indexing_fallback_on_getitem_3d(store: StorePath, index, exp (([0, 2], slice(None, None, 2)), [[1, 0, 1], [0, 0, 0], [1, 0, 1]]), ], ) -def test_orthogonal_indexing_fallback_on_setitem_2d(store: StorePath, index, expected_result): +def test_orthogonal_indexing_fallback_on_setitem_2d( + store: StorePath, index: BasicSelection, expected_result: list[list[int]] +) -> None: """ Tests the orthogonal indexing fallback on __setitem__ for a 3D matrix. @@ -468,7 +483,7 @@ def test_orthogonal_indexing_fallback_on_setitem_2d(store: StorePath, index, exp np.testing.assert_array_equal(z[:], a, err_msg="Indexing disagrees with numpy") -def test_fancy_indexing_doesnt_mix_with_implicit_slicing(store: StorePath): +def test_fancy_indexing_doesnt_mix_with_implicit_slicing(store: StorePath) -> None: z2 = zarr_array_from_numpy_array(store, np.zeros((5, 5, 5))) with pytest.raises(IndexError): z2[[1, 2, 3], [1, 2, 3]] = 2 @@ -522,7 +537,7 @@ def test_set_basic_selection_0d( # arr_z[..., "foo", "bar"] = v[["foo", "bar"]] -def _test_get_orthogonal_selection(a, z, selection): +def _test_get_orthogonal_selection(a: Array, z: Array, selection: BasicSelection) -> None: expect = oindex(a, selection) actual = z.get_orthogonal_selection(selection) assert_array_equal(expect, actual) @@ -531,7 +546,7 @@ def _test_get_orthogonal_selection(a, z, selection): # noinspection PyStatementEffect -def test_get_orthogonal_selection_1d_bool(store: StorePath): +def test_get_orthogonal_selection_1d_bool(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -552,7 +567,7 @@ def test_get_orthogonal_selection_1d_bool(store: StorePath): # noinspection PyStatementEffect -def test_get_orthogonal_selection_1d_int(store: StorePath): +def test_get_orthogonal_selection_1d_int(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -591,7 +606,9 @@ def test_get_orthogonal_selection_1d_int(store: StorePath): z.oindex[selection] -def _test_get_orthogonal_selection_2d(a, z, ix0, ix1): +def _test_get_orthogonal_selection_2d( + a: Array, z: Array, ix0: BasicSelection, ix1: BasicSelection +) -> None: selections = [ # index both axes with array (ix0, ix1), @@ -609,7 +626,7 @@ def _test_get_orthogonal_selection_2d(a, z, ix0, ix1): # noinspection PyStatementEffect -def test_get_orthogonal_selection_2d(store: StorePath): +def test_get_orthogonal_selection_2d(store: StorePath) -> None: # setup a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) @@ -651,7 +668,9 @@ def test_get_orthogonal_selection_2d(store: StorePath): z.oindex[selection] -def _test_get_orthogonal_selection_3d(a, z, ix0, ix1, ix2): +def _test_get_orthogonal_selection_3d( + a: Array, z: Array, ix0: BasicSelection, ix1: BasicSelection, ix2: BasicSelection +) -> None: selections = [ # single value (84, 42, 4), @@ -686,7 +705,7 @@ def _test_get_orthogonal_selection_3d(a, z, ix0, ix1, ix2): _test_get_orthogonal_selection(a, z, selection) -def test_get_orthogonal_selection_3d(store: StorePath): +def test_get_orthogonal_selection_3d(store: StorePath) -> None: # setup a = np.arange(100000, dtype=int).reshape(200, 50, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(60, 20, 3)) @@ -715,7 +734,7 @@ def test_get_orthogonal_selection_3d(store: StorePath): _test_get_orthogonal_selection_3d(a, z, ix0, ix1, ix2) -def test_orthogonal_indexing_edge_cases(store: StorePath): +def test_orthogonal_indexing_edge_cases(store: StorePath) -> None: a = np.arange(6).reshape(1, 2, 3) z = zarr_array_from_numpy_array(store, a, chunk_shape=(1, 2, 3)) @@ -728,7 +747,7 @@ def test_orthogonal_indexing_edge_cases(store: StorePath): assert_array_equal(expect, actual) -def _test_set_orthogonal_selection(v, a, z, selection): +def _test_set_orthogonal_selection(v: Array, a: Array, z: Array, selection: BasicSelection) -> None: for value in 42, oindex(v, selection), oindex(v, selection).tolist(): if isinstance(value, list) and value == []: # skip these cases as cannot preserve all dimensions @@ -746,7 +765,7 @@ def _test_set_orthogonal_selection(v, a, z, selection): assert_array_equal(a, z[:]) -def test_set_orthogonal_selection_1d(store: StorePath): +def test_set_orthogonal_selection_1d(store: StorePath) -> None: # setup v = np.arange(1050, dtype=int) a = np.empty(v.shape, dtype=int) @@ -772,7 +791,9 @@ def test_set_orthogonal_selection_1d(store: StorePath): _test_set_orthogonal_selection(v, a, z, selection) -def _test_set_orthogonal_selection_2d(v, a, z, ix0, ix1): +def _test_set_orthogonal_selection_2d( + v: Array, a: Array, z: Array, ix0: BasicSelection, ix1: BasicSelection +) -> None: selections = [ # index both axes with array (ix0, ix1), @@ -786,7 +807,7 @@ def _test_set_orthogonal_selection_2d(v, a, z, ix0, ix1): _test_set_orthogonal_selection(v, a, z, selection) -def test_set_orthogonal_selection_2d(store: StorePath): +def test_set_orthogonal_selection_2d(store: StorePath) -> None: # setup v = np.arange(10000, dtype=int).reshape(1000, 10) a = np.empty_like(v) @@ -815,7 +836,9 @@ def test_set_orthogonal_selection_2d(store: StorePath): _test_set_orthogonal_selection(v, a, z, selection) -def _test_set_orthogonal_selection_3d(v, a, z, ix0, ix1, ix2): +def _test_set_orthogonal_selection_3d( + v: Array, a: Array, z: Array, ix0: BasicSelection, ix1: BasicSelection, ix2: BasicSelection +) -> None: selections = ( # single value (84, 42, 4), @@ -841,7 +864,7 @@ def _test_set_orthogonal_selection_3d(v, a, z, ix0, ix1, ix2): _test_set_orthogonal_selection(v, a, z, selection) -def test_set_orthogonal_selection_3d(store: StorePath): +def test_set_orthogonal_selection_3d(store: StorePath) -> None: # setup v = np.arange(100000, dtype=int).reshape(200, 50, 10) a = np.empty_like(v) @@ -875,7 +898,7 @@ def test_set_orthogonal_selection_3d(store: StorePath): _test_set_orthogonal_selection_3d(v, a, z, ix0, ix1, ix2) -def test_orthogonal_indexing_fallback_on_get_setitem(store: StorePath): +def test_orthogonal_indexing_fallback_on_get_setitem(store: StorePath) -> None: z = zarr_array_from_numpy_array(store, np.zeros((20, 20))) z[[1, 2, 3], [1, 2, 3]] = 1 np.testing.assert_array_equal( @@ -896,7 +919,7 @@ def test_orthogonal_indexing_fallback_on_get_setitem(store: StorePath): np.testing.assert_array_equal(z2[:], [0, 1, 1, 1, 0]) -def _test_get_coordinate_selection(a, z, selection): +def _test_get_coordinate_selection(a: Array, z: Array, selection: BasicSelection) -> None: expect = a[selection] actual = z.get_coordinate_selection(selection) assert_array_equal(expect, actual) @@ -920,7 +943,7 @@ def _test_get_coordinate_selection(a, z, selection): # noinspection PyStatementEffect -def test_get_coordinate_selection_1d(store: StorePath): +def test_get_coordinate_selection_1d(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -962,7 +985,7 @@ def test_get_coordinate_selection_1d(store: StorePath): z.vindex[selection] -def test_get_coordinate_selection_2d(store: StorePath): +def test_get_coordinate_selection_2d(store: StorePath) -> None: # setup a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) @@ -1016,7 +1039,7 @@ def test_get_coordinate_selection_2d(store: StorePath): z.get_coordinate_selection(selection) -def _test_set_coordinate_selection(v, a, z, selection): +def _test_set_coordinate_selection(v: Array, a: Array, z: Array, selection: BasicSelection) -> None: for value in 42, v[selection], v[selection].tolist(): # setup expectation a[:] = 0 @@ -1031,7 +1054,7 @@ def _test_set_coordinate_selection(v, a, z, selection): assert_array_equal(a, z[:]) -def test_set_coordinate_selection_1d(store: StorePath): +def test_set_coordinate_selection_1d(store: StorePath) -> None: # setup v = np.arange(1050, dtype=int) a = np.empty(v.shape, dtype=v.dtype) @@ -1055,7 +1078,7 @@ def test_set_coordinate_selection_1d(store: StorePath): z.vindex[selection] = 42 -def test_set_coordinate_selection_2d(store: StorePath): +def test_set_coordinate_selection_2d(store: StorePath) -> None: # setup v = np.arange(10000, dtype=int).reshape(1000, 10) a = np.empty_like(v) @@ -1086,7 +1109,9 @@ def test_set_coordinate_selection_2d(store: StorePath): _test_set_coordinate_selection(v, a, z, (ix0, ix1)) -def _test_get_block_selection(a, z, selection, expected_idx): +def _test_get_block_selection( + a: Array, z: Array, selection: BasicSelection, expected_idx: BasicSelection +) -> None: expect = a[expected_idx] actual = z.get_block_selection(selection) assert_array_equal(expect, actual) @@ -1138,7 +1163,7 @@ def _test_get_block_selection(a, z, selection, expected_idx): ] -def test_get_block_selection_1d(store: StorePath): +def test_get_block_selection_1d(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -1191,7 +1216,7 @@ def test_get_block_selection_1d(store: StorePath): ] -def test_get_block_selection_2d(store: StorePath): +def test_get_block_selection_2d(store: StorePath) -> None: # setup a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) @@ -1212,7 +1237,13 @@ def test_get_block_selection_2d(store: StorePath): z.get_block_selection(selection) -def _test_set_block_selection(v: np.ndarray, a: np.ndarray, z: zarr.Array, selection, expected_idx): +def _test_set_block_selection( + v: np.ndarray, + a: np.ndarray, + z: zarr.Array, + selection: BasicSelection, + expected_idx: BasicSelection, +) -> None: for value in 42, v[expected_idx], v[expected_idx].tolist(): # setup expectation a[:] = 0 @@ -1227,7 +1258,7 @@ def _test_set_block_selection(v: np.ndarray, a: np.ndarray, z: zarr.Array, selec assert_array_equal(a, z[:]) -def test_set_block_selection_1d(store: StorePath): +def test_set_block_selection_1d(store: StorePath) -> None: # setup v = np.arange(1050, dtype=int) a = np.empty(v.shape, dtype=v.dtype) @@ -1245,7 +1276,7 @@ def test_set_block_selection_1d(store: StorePath): z.blocks[selection] = 42 -def test_set_block_selection_2d(store: StorePath): +def test_set_block_selection_2d(store: StorePath) -> None: # setup v = np.arange(10000, dtype=int).reshape(1000, 10) a = np.empty(v.shape, dtype=v.dtype) @@ -1267,7 +1298,7 @@ def test_set_block_selection_2d(store: StorePath): z.set_block_selection(selection, 42) -def _test_get_mask_selection(a, z, selection): +def _test_get_mask_selection(a: Array, z: Array, selection: BasicSelection) -> None: expect = a[selection] actual = z.get_mask_selection(selection) assert_array_equal(expect, actual) @@ -1293,7 +1324,7 @@ def _test_get_mask_selection(a, z, selection): # noinspection PyStatementEffect -def test_get_mask_selection_1d(store: StorePath): +def test_get_mask_selection_1d(store: StorePath) -> None: # setup a = np.arange(1050, dtype=int) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -1318,7 +1349,7 @@ def test_get_mask_selection_1d(store: StorePath): # noinspection PyStatementEffect -def test_get_mask_selection_2d(store: StorePath): +def test_get_mask_selection_2d(store: StorePath) -> None: # setup a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) @@ -1338,7 +1369,7 @@ def test_get_mask_selection_2d(store: StorePath): z.vindex[[True, False]] # wrong no. dimensions -def _test_set_mask_selection(v, a, z, selection): +def _test_set_mask_selection(v: Array, a: Array, z: Array, selection: BasicSelection) -> None: a[:] = 0 z[:] = 0 a[selection] = v[selection] @@ -1352,7 +1383,7 @@ def _test_set_mask_selection(v, a, z, selection): assert_array_equal(a, z[:]) -def test_set_mask_selection_1d(store: StorePath): +def test_set_mask_selection_1d(store: StorePath) -> None: # setup v = np.arange(1050, dtype=int) a = np.empty_like(v) @@ -1371,7 +1402,7 @@ def test_set_mask_selection_1d(store: StorePath): z.vindex[selection] = 42 -def test_set_mask_selection_2d(store: StorePath): +def test_set_mask_selection_2d(store: StorePath) -> None: # setup v = np.arange(10000, dtype=int).reshape(1000, 10) a = np.empty_like(v) @@ -1384,7 +1415,7 @@ def test_set_mask_selection_2d(store: StorePath): _test_set_mask_selection(v, a, z, ix) -def test_get_selection_out(store: StorePath): +def test_get_selection_out(store: StorePath) -> None: # basic selections a = np.arange(1050) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) @@ -1454,7 +1485,7 @@ def test_get_selection_out(store: StorePath): @pytest.mark.xfail(reason="fields are not supported in v3") -def test_get_selections_with_fields(store: StorePath): +def test_get_selections_with_fields(store: StorePath) -> None: a = [("aaa", 1, 4.2), ("bbb", 2, 8.4), ("ccc", 3, 12.6)] a = np.array(a, dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")]) z = zarr_array_from_numpy_array(store, a, chunk_shape=(2,)) @@ -1560,7 +1591,7 @@ def test_get_selections_with_fields(store: StorePath): @pytest.mark.xfail(reason="fields are not supported in v3") -def test_set_selections_with_fields(store: StorePath): +def test_set_selections_with_fields(store: StorePath) -> None: v = [("aaa", 1, 4.2), ("bbb", 2, 8.4), ("ccc", 3, 12.6)] v = np.array(v, dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")]) a = np.empty_like(v) @@ -1643,14 +1674,14 @@ def test_set_selections_with_fields(store: StorePath): assert_array_equal(a, z[:]) -def test_slice_selection_uints(): +def test_slice_selection_uints() -> None: arr = np.arange(24).reshape((4, 6)) idx = np.uint64(3) slice_sel = make_slice_selection((idx,)) assert arr[tuple(slice_sel)].shape == (1, 6) -def test_numpy_int_indexing(store: StorePath): +def test_numpy_int_indexing(store: StorePath) -> None: a = np.arange(1050) z = zarr_array_from_numpy_array(store, a, chunk_shape=(100,)) assert a[42] == z[42] @@ -1682,7 +1713,9 @@ def test_numpy_int_indexing(store: StorePath): ), ], ) -async def test_accessed_chunks(shape, chunks, ops): +async def test_accessed_chunks( + shape: ChunkCoords, chunks: ChunkCoords, ops: tuple[str, tuple[int, ...]] +) -> None: # Test that only the required chunks are accessed during basic selection operations # shape: array shape # chunks: chunk size @@ -1755,7 +1788,7 @@ async def test_accessed_chunks(shape, chunks, ops): [[100, 200, 300], [4, 5, 6]], ], ) -def test_indexing_equals_numpy(store, selection): +def test_indexing_equals_numpy(store: MemoryStore, selection: BasicSelection) -> None: a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) # note: in python 3.10 a[*selection] is not valid unpacking syntax @@ -1773,7 +1806,9 @@ def test_indexing_equals_numpy(store, selection): [np.full(1000, True), [True, False] * 5], ], ) -def test_orthogonal_bool_indexing_like_numpy_ix(store, selection): +def test_orthogonal_bool_indexing_like_numpy_ix( + store: MemoryStore, selection: BasicSelection +) -> None: a = np.arange(10000, dtype=int).reshape(1000, 10) z = zarr_array_from_numpy_array(store, a, chunk_shape=(300, 3)) expected = a[np.ix_(*selection)] From 8dceef2bcd5f385242db4703da5c7ead7269a29d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 6 Aug 2024 12:15:28 +0200 Subject: [PATCH 2/4] add docstring with explanation to new test --- tests/v3/test_indexing.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/v3/test_indexing.py b/tests/v3/test_indexing.py index c84cc3eb24..e170852e9a 100644 --- a/tests/v3/test_indexing.py +++ b/tests/v3/test_indexing.py @@ -15,6 +15,7 @@ from zarr.array import Array from zarr.buffer import Buffer, BufferPrototype, NDBuffer from zarr.common import ChunkCoords +from zarr.group import Group from zarr.indexing import ( BasicSelection, make_slice_selection, @@ -1815,3 +1816,26 @@ def test_orthogonal_bool_indexing_like_numpy_ix( # note: in python 3.10 z[*selection] is not valid unpacking syntax actual = z[(*selection,)] assert_array_equal(expected, actual, err_msg=f"{selection=}") + + +def test_indexing_after_close(store: StorePath) -> None: + """ + Test that data is persisted after calling store.close, and subsequently invoking array[:]. + Bug discovered here: https://github.com/zarr-developers/zarr-python/pull/1746#issuecomment-2269562717 + """ + root = Group.create(store) + value = 1 + nparray = np.array([value], dtype=np.int8) + + a = root.create_array( + "/0/0", + shape=nparray.shape, + chunks=(1,), + dtype=nparray.dtype.str, + attributes={}, + fill_value=nparray.dtype.type(0), + ) + a[:] = nparray + assert a[:] == value + store.store.close() + assert a[:] == value From acdb9ef02990d0cd26547fd37d970a244c52f865 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Tue, 6 Aug 2024 15:16:00 +0200 Subject: [PATCH 3/4] raise RuntimeError when attempting to set or get a closed store --- src/zarr/abc/store.py | 4 ---- src/zarr/store/core.py | 1 - src/zarr/store/local.py | 2 +- src/zarr/store/memory.py | 6 +++--- src/zarr/store/remote.py | 3 ++- tests/v3/test_buffer.py | 14 ++++++++------ tests/v3/test_config.py | 11 ++++++----- tests/v3/test_indexing.py | 6 ++++-- tests/v3/test_store/test_memory.py | 4 ++-- 9 files changed, 26 insertions(+), 25 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index 449816209b..ad1514289c 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -52,10 +52,6 @@ async def _open(self) -> None: raise FileExistsError("Store already exists") self._is_open = True - async def _ensure_open(self) -> None: - if not self._is_open: - await self._open() - @abstractmethod async def empty(self) -> bool: ... diff --git a/src/zarr/store/core.py b/src/zarr/store/core.py index fa35879308..2aabe61a49 100644 --- a/src/zarr/store/core.py +++ b/src/zarr/store/core.py @@ -80,7 +80,6 @@ async def make_store_path( elif isinstance(store_like, Store): if mode is not None: assert AccessMode.from_literal(mode) == store_like.mode - await store_like._ensure_open() return StorePath(store_like) elif store_like is None: if mode is None: diff --git a/src/zarr/store/local.py b/src/zarr/store/local.py index 25fd9fc13a..1e26de9e07 100644 --- a/src/zarr/store/local.py +++ b/src/zarr/store/local.py @@ -73,11 +73,11 @@ class LocalStore(Store): root: Path def __init__(self, root: Path | str, *, mode: AccessModeLiteral = "r"): - super().__init__(mode=mode) if isinstance(root, str): root = Path(root) assert isinstance(root, Path) self.root = root + super().__init__(mode=mode) async def clear(self) -> None: self._check_writable() diff --git a/src/zarr/store/memory.py b/src/zarr/store/memory.py index dd3e52e703..25fdbd41d4 100644 --- a/src/zarr/store/memory.py +++ b/src/zarr/store/memory.py @@ -23,8 +23,8 @@ def __init__( *, mode: AccessModeLiteral = "r", ): - super().__init__(mode=mode) self._store_dict = store_dict or {} + super().__init__(mode=mode) async def empty(self) -> bool: return not self._store_dict @@ -45,7 +45,7 @@ async def get( byte_range: tuple[int | None, int | None] | None = None, ) -> Buffer | None: if not self._is_open: - await self._open() + raise RuntimeError("Store is closed. Cannot `get` from a closed store.") assert isinstance(key, str) try: value = self._store_dict[key] @@ -71,7 +71,7 @@ async def exists(self, key: str) -> bool: async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: if not self._is_open: - await self._open() + raise RuntimeError("Store is closed. Cannot `get` from a closed store.") self._check_writable() assert isinstance(key, str) if not isinstance(value, Buffer): diff --git a/src/zarr/store/remote.py b/src/zarr/store/remote.py index c742d9e567..83361afcd8 100644 --- a/src/zarr/store/remote.py +++ b/src/zarr/store/remote.py @@ -50,7 +50,7 @@ def __init__( storage_options: passed on to fsspec to make the filesystem instance. If url is a UPath, this must not be used. """ - super().__init__(mode=mode) + if isinstance(url, str): self._url = url.rstrip("/") self._fs, _path = fsspec.url_to_fs(url, **storage_options) @@ -73,6 +73,7 @@ def __init__( # test instantiate file system if not self._fs.async_impl: raise TypeError("FileSystem needs to support async operations") + super().__init__(mode=mode) async def clear(self) -> None: try: diff --git a/tests/v3/test_buffer.py b/tests/v3/test_buffer.py index d53e98d42d..1b72aaacb6 100644 --- a/tests/v3/test_buffer.py +++ b/tests/v3/test_buffer.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Any + import numpy as np import pytest @@ -20,19 +22,19 @@ ) -def test_nd_array_like(xp): +def test_nd_array_like(xp: Any) -> None: ary = xp.arange(10) assert isinstance(ary, ArrayLike) assert isinstance(ary, NDArrayLike) @pytest.mark.asyncio -async def test_async_array_prototype(): +async def test_async_array_prototype() -> None: """Test the use of a custom buffer prototype""" expect = np.zeros((9, 9), dtype="uint16", order="F") a = await AsyncArray.create( - StorePath(StoreExpectingTestBuffer(mode="w")) / "test_async_array_prototype", + StorePath(await StoreExpectingTestBuffer.open(mode="w")) / "test_async_array_prototype", shape=expect.shape, chunk_shape=(5, 5), dtype=expect.dtype, @@ -53,10 +55,10 @@ async def test_async_array_prototype(): @pytest.mark.asyncio -async def test_codecs_use_of_prototype(): +async def test_codecs_use_of_prototype() -> None: expect = np.zeros((10, 10), dtype="uint16", order="F") a = await AsyncArray.create( - StorePath(StoreExpectingTestBuffer(mode="w")) / "test_codecs_use_of_prototype", + StorePath(await StoreExpectingTestBuffer.open(mode="w")) / "test_codecs_use_of_prototype", shape=expect.shape, chunk_shape=(5, 5), dtype=expect.dtype, @@ -84,7 +86,7 @@ async def test_codecs_use_of_prototype(): assert np.array_equal(expect, got) -def test_numpy_buffer_prototype(): +def test_numpy_buffer_prototype() -> None: buffer = numpy_buffer_prototype().buffer.create_zero_length() ndbuffer = numpy_buffer_prototype().nd_buffer.create(shape=(1, 2), dtype=np.dtype("int64")) assert isinstance(buffer.as_array_like(), np.ndarray) diff --git a/tests/v3/test_config.py b/tests/v3/test_config.py index 8e7b868520..956db79df2 100644 --- a/tests/v3/test_config.py +++ b/tests/v3/test_config.py @@ -27,6 +27,7 @@ register_ndbuffer, register_pipeline, ) +from zarr.sync import sync from zarr.testing.buffer import ( NDBufferUsingTestNDArrayLike, StoreExpectingTestBuffer, @@ -191,11 +192,11 @@ def test_config_ndbuffer_implementation(store): assert isinstance(got, TestNDArrayLike) -def test_config_buffer_implementation(): +def test_config_buffer_implementation() -> None: # has default value assert fully_qualified_name(get_buffer_class()) == config.defaults[0]["buffer"] - - arr = zeros(shape=(100), store=StoreExpectingTestBuffer(mode="w")) + store = sync(StoreExpectingTestBuffer.open(mode="w")) + arr = zeros(shape=(100), store=store) # AssertionError of StoreExpectingTestBuffer when not using my buffer with pytest.raises(AssertionError): @@ -213,7 +214,7 @@ def test_config_buffer_implementation(): data2d = np.arange(1000).reshape(100, 10) arr_sharding = zeros( shape=(100, 10), - store=StoreExpectingTestBuffer(mode="w"), + store=sync(StoreExpectingTestBuffer.open(mode="w")), codecs=[ShardingCodec(chunk_shape=(10, 10))], ) arr_sharding[:] = data2d @@ -221,7 +222,7 @@ def test_config_buffer_implementation(): arr_Crc32c = zeros( shape=(100, 10), - store=StoreExpectingTestBuffer(mode="w"), + store=sync(StoreExpectingTestBuffer.open(mode="w")), codecs=[BytesCodec(), Crc32cCodec()], ) arr_Crc32c[:] = data2d diff --git a/tests/v3/test_indexing.py b/tests/v3/test_indexing.py index e170852e9a..f913d824ab 100644 --- a/tests/v3/test_indexing.py +++ b/tests/v3/test_indexing.py @@ -27,11 +27,12 @@ from zarr.registry import get_ndbuffer_class from zarr.store.core import StorePath from zarr.store.memory import MemoryStore +from zarr.sync import sync @pytest.fixture def store() -> Iterator[StorePath]: - yield StorePath(MemoryStore(mode="w")) + yield StorePath(sync(MemoryStore.open(mode="w"))) def zarr_array_from_numpy_array( @@ -1838,4 +1839,5 @@ def test_indexing_after_close(store: StorePath) -> None: a[:] = nparray assert a[:] == value store.store.close() - assert a[:] == value + with pytest.raises(RuntimeError, match="Store is closed"): + assert a[:] == value diff --git a/tests/v3/test_store/test_memory.py b/tests/v3/test_store/test_memory.py index 5b8f1ef875..63a631020a 100644 --- a/tests/v3/test_store/test_memory.py +++ b/tests/v3/test_store/test_memory.py @@ -23,8 +23,8 @@ def store_kwargs( return {"store_dict": request.param, "mode": "r+"} @pytest.fixture(scope="function") - def store(self, store_kwargs: str | None | dict[str, Buffer]) -> MemoryStore: - return self.store_cls(**store_kwargs) + async def store(self, store_kwargs: str | None | dict[str, Buffer]) -> MemoryStore: + return await self.store_cls.open(**store_kwargs) def test_store_repr(self, store: MemoryStore) -> None: assert str(store) == f"memory://{id(store._store_dict)}" From 17ee300cc41092323df7ad9d421f74a43ea9b4fe Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Mon, 12 Aug 2024 13:18:41 +0200 Subject: [PATCH 4/4] Update src/zarr/store/memory.py Co-authored-by: Hannes Spitz <44113112+brokkoli71@users.noreply.github.com> --- src/zarr/store/memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zarr/store/memory.py b/src/zarr/store/memory.py index 25fdbd41d4..981becaa9b 100644 --- a/src/zarr/store/memory.py +++ b/src/zarr/store/memory.py @@ -71,7 +71,7 @@ async def exists(self, key: str) -> bool: async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: if not self._is_open: - raise RuntimeError("Store is closed. Cannot `get` from a closed store.") + raise RuntimeError("Store is closed. Cannot `set` in a closed store.") self._check_writable() assert isinstance(key, str) if not isinstance(value, Buffer):