Merge branch 'v3' of https://github.com/zarr-developers/zarr-python into fix/dask-compat

Author: jhamman

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_language_version:
   python: python3
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.4
+    rev: v0.6.5
     hooks:
     - id: ruff
       args: ["--fix", "--show-fixes"]

src/zarr/core/buffer/core.py

@@ -462,8 +462,12 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return f"<NDBuffer shape={self.shape} dtype={self.dtype} {self._data!r}>"
 
-    def all_equal(self, other: Any) -> bool:
-        return bool((self._data == other).all())
+    def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
+        """Compare to `other` using np.array_equal."""
+        # use array_equal to obtain equal_nan=True functionality
+        data, other = np.broadcast_arrays(self._data, other)
+        result = np.array_equal(self._data, other, equal_nan=equal_nan)
+        return result
 
     def fill(self, value: Any) -> None:
         self._data.fill(value)

src/zarr/testing/store.py

@@ -251,10 +251,10 @@ async def test_list_dir(self, store: S) -> None:
     async def test_set_get(self, store_kwargs: dict[str, Any]) -> None:
         kwargs = {**store_kwargs, **{"mode": "w"}}
         store = self.store_cls(**kwargs)
-        await zarr.api.asynchronous.open_array(store=store, path="a", mode="w", shape=(4,))
+        await zarr.api.asynchronous.open_array(store=store, path="a", shape=(4,))
         keys = [x async for x in store.list()]
         assert keys == ["a/zarr.json"]
 
         # no errors
-        await zarr.api.asynchronous.open_array(store=store, path="a", mode="r")
-        await zarr.api.asynchronous.open_array(store=store, path="a", mode="a")
+        await zarr.api.asynchronous.open_array(store=store, path="a")
+        await zarr.api.asynchronous.open_array(store=store, path="a")

tests/v3/test_array.py

@@ -138,6 +138,25 @@ def test_array_v3_fill_value(store: MemoryStore, fill_value: int, dtype_str: str
     assert arr.fill_value.dtype == arr.dtype
 
 
+@pytest.mark.parametrize("store", ["memory"], indirect=True)
+async def test_array_v3_nan_fill_value(store: MemoryStore) -> None:
+    shape = (10,)
+    arr = Array.create(
+        store=store,
+        shape=shape,
+        dtype=np.float64,
+        zarr_format=3,
+        chunk_shape=shape,
+        fill_value=np.nan,
+    )
+    arr[:] = np.nan
+
+    assert np.isnan(arr.fill_value)
+    assert arr.fill_value.dtype == arr.dtype
+    # all fill value chunk is an empty chunk, and should not be written
+    assert len([a async for a in store.list_prefix("/")]) == 0
+
+
 @pytest.mark.parametrize("store", ("local",), indirect=["store"])
 @pytest.mark.parametrize("zarr_format", (2, 3))
 async def test_serializable_async_array(

tests/v3/test_group.py

@@ -6,14 +6,15 @@
 import numpy as np
 import pytest
 
+import zarr
 from zarr import Array, AsyncArray, AsyncGroup, Group
 from zarr.abc.store import Store
 from zarr.core.buffer import default_buffer_prototype
 from zarr.core.common import JSON, ZarrFormat
 from zarr.core.group import GroupMetadata
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
-from zarr.storage import LocalStore, StorePath
+from zarr.storage import LocalStore, MemoryStore, StorePath
 from zarr.storage.common import make_store_path
 
 from .conftest import parse_store
@@ -699,3 +700,154 @@ def test_serializable_sync_group(store: LocalStore, zarr_format: ZarrFormat) ->
     actual = pickle.loads(p)
 
     assert actual == expected
+
+
+async def test_group_members_async(store: LocalStore | MemoryStore) -> None:
+    group = AsyncGroup(
+        GroupMetadata(),
+        store_path=StorePath(store=store, path="root"),
+    )
+    a0 = await group.create_array("a0", shape=(1,))
+    g0 = await group.create_group("g0")
+    a1 = await g0.create_array("a1", shape=(1,))
+    g1 = await g0.create_group("g1")
+    a2 = await g1.create_array("a2", shape=(1,))
+    g2 = await g1.create_group("g2")
+
+    # immediate children
+    children = sorted([x async for x in group.members()], key=lambda x: x[0])
+    assert children == [
+        ("a0", a0),
+        ("g0", g0),
+    ]
+
+    nmembers = await group.nmembers()
+    assert nmembers == 2
+
+    # partial
+    children = sorted([x async for x in group.members(max_depth=1)], key=lambda x: x[0])
+    expected = [
+        ("a0", a0),
+        ("g0", g0),
+        ("g0/a1", a1),
+        ("g0/g1", g1),
+    ]
+    assert children == expected
+    nmembers = await group.nmembers(max_depth=1)
+    assert nmembers == 4
+
+    # all children
+    all_children = sorted([x async for x in group.members(max_depth=None)], key=lambda x: x[0])
+    expected = [
+        ("a0", a0),
+        ("g0", g0),
+        ("g0/a1", a1),
+        ("g0/g1", g1),
+        ("g0/g1/a2", a2),
+        ("g0/g1/g2", g2),
+    ]
+    assert all_children == expected
+
+    nmembers = await group.nmembers(max_depth=None)
+    assert nmembers == 6
+
+    with pytest.raises(ValueError, match="max_depth"):
+        [x async for x in group.members(max_depth=-1)]
+
+
+async def test_require_group(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
+    root = await AsyncGroup.create(store=store, zarr_format=zarr_format)
+
+    # create foo group
+    _ = await root.create_group("foo", attributes={"foo": 100})
+
+    # test that we can get the group using require_group
+    foo_group = await root.require_group("foo")
+    assert foo_group.attrs == {"foo": 100}
+
+    # test that we can get the group using require_group and overwrite=True
+    foo_group = await root.require_group("foo", overwrite=True)
+
+    _ = await foo_group.create_array(
+        "bar", shape=(10,), dtype="uint8", chunk_shape=(2,), attributes={"foo": 100}
+    )
+
+    # test that overwriting a group w/ children fails
+    # TODO: figure out why ensure_no_existing_node is not catching the foo.bar array
+    #
+    # with pytest.raises(ContainsArrayError):
+    #     await root.require_group("foo", overwrite=True)
+
+    # test that requiring a group where an array is fails
+    with pytest.raises(TypeError):
+        await foo_group.require_group("bar")
+
+
+async def test_require_groups(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
+    root = await AsyncGroup.create(store=store, zarr_format=zarr_format)
+    # create foo group
+    _ = await root.create_group("foo", attributes={"foo": 100})
+    # create bar group
+    _ = await root.create_group("bar", attributes={"bar": 200})
+
+    foo_group, bar_group = await root.require_groups("foo", "bar")
+    assert foo_group.attrs == {"foo": 100}
+    assert bar_group.attrs == {"bar": 200}
+
+    # get a mix of existing and new groups
+    foo_group, spam_group = await root.require_groups("foo", "spam")
+    assert foo_group.attrs == {"foo": 100}
+    assert spam_group.attrs == {}
+
+    # no names
+    no_group = await root.require_groups()
+    assert no_group == ()
+
+
+async def test_create_dataset(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
+    root = await AsyncGroup.create(store=store, zarr_format=zarr_format)
+    with pytest.warns(DeprecationWarning):
+        foo = await root.create_dataset("foo", shape=(10,), dtype="uint8")
+    assert foo.shape == (10,)
+
+    with pytest.raises(ContainsArrayError), pytest.warns(DeprecationWarning):
+        await root.create_dataset("foo", shape=(100,), dtype="int8")
+
+    _ = await root.create_group("bar")
+    with pytest.raises(ContainsGroupError), pytest.warns(DeprecationWarning):
+        await root.create_dataset("bar", shape=(100,), dtype="int8")
+
+
+async def test_require_array(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
+    root = await AsyncGroup.create(store=store, zarr_format=zarr_format)
+    foo1 = await root.require_array("foo", shape=(10,), dtype="i8", attributes={"foo": 101})
+    assert foo1.attrs == {"foo": 101}
+    foo2 = await root.require_array("foo", shape=(10,), dtype="i8")
+    assert foo2.attrs == {"foo": 101}
+
+    # exact = False
+    _ = await root.require_array("foo", shape=10, dtype="f8")
+
+    # errors w/ exact True
+    with pytest.raises(TypeError, match="Incompatible dtype"):
+        await root.require_array("foo", shape=(10,), dtype="f8", exact=True)
+
+    with pytest.raises(TypeError, match="Incompatible shape"):
+        await root.require_array("foo", shape=(100, 100), dtype="i8")
+
+    with pytest.raises(TypeError, match="Incompatible dtype"):
+        await root.require_array("foo", shape=(10,), dtype="f4")
+
+    _ = await root.create_group("bar")
+    with pytest.raises(TypeError, match="Incompatible object"):
+        await root.require_array("bar", shape=(10,), dtype="int8")
+
+
+async def test_open_mutable_mapping():
+    group = await zarr.api.asynchronous.open_group(store={}, mode="w")
+    assert isinstance(group.store_path.store, MemoryStore)
+
+
+def test_open_mutable_mapping_sync():
+    group = zarr.open_group(store={}, mode="w")
+    assert isinstance(group.store_path.store, MemoryStore)

tests/v3/test_properties.py

@@ -18,11 +18,6 @@ def test_roundtrip(data: st.DataObject) -> None:
 
 
 @given(data=st.data())
-# The filter warning here is to silence an occasional warning in NDBuffer.all_equal
-# See https://github.com/zarr-developers/zarr-python/pull/2118#issuecomment-2310280899
-# Uncomment the next line to reproduce the original failure.
-# @reproduce_failure('6.111.2', b'AXicY2FgZGRAB/8/ndR2z7nkDZEDADWpBL4=')
-@pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_basic_indexing(data: st.DataObject) -> None:
     zarray = data.draw(arrays())
     nparray = zarray[:]
@@ -37,11 +32,6 @@ def test_basic_indexing(data: st.DataObject) -> None:
 
 
 @given(data=st.data())
-# The filter warning here is to silence an occasional warning in NDBuffer.all_equal
-# See https://github.com/zarr-developers/zarr-python/pull/2118#issuecomment-2310280899
-# Uncomment the next line to reproduce the original failure.
-# @reproduce_failure('6.111.2', b'AXicY2FgZGRAB/8/eLmF7qr/C5EDADZUBRM=')
-@pytest.mark.filterwarnings("ignore::RuntimeWarning")
 def test_vindex(data: st.DataObject) -> None:
     zarray = data.draw(arrays())
     nparray = zarray[:]

tests/v3/test_store/test_memory.py

@@ -4,9 +4,10 @@
 
 import pytest
 
-from zarr.core.buffer import Buffer, cpu
-from zarr.storage.memory import MemoryStore
+from zarr.core.buffer import Buffer, cpu, gpu
+from zarr.storage.memory import GpuMemoryStore, MemoryStore
 from zarr.testing.store import StoreTests
+from zarr.testing.utils import gpu_test
 
 
 class TestMemoryStore(StoreTests[MemoryStore, cpu.Buffer]):
@@ -56,3 +57,48 @@ def test_serizalizable_store(self, store: MemoryStore) -> None:
 
         with pytest.raises(NotImplementedError):
             pickle.dumps(store)
+
+
+@gpu_test
+class TestGpuMemoryStore(StoreTests[GpuMemoryStore, gpu.Buffer]):
+    store_cls = GpuMemoryStore
+    buffer_cls = gpu.Buffer
+
+    def set(self, store: GpuMemoryStore, key: str, value: Buffer) -> None:
+        store._store_dict[key] = value
+
+    def get(self, store: MemoryStore, key: str) -> Buffer:
+        return store._store_dict[key]
+
+    @pytest.fixture(scope="function", params=[None, {}])
+    def store_kwargs(self, request) -> dict[str, str | None | dict[str, Buffer]]:
+        return {"store_dict": request.param, "mode": "r+"}
+
+    @pytest.fixture(scope="function")
+    def store(self, store_kwargs: str | None | dict[str, gpu.Buffer]) -> GpuMemoryStore:
+        return self.store_cls(**store_kwargs)
+
+    def test_store_repr(self, store: GpuMemoryStore) -> None:
+        assert str(store) == f"gpumemory://{id(store._store_dict)}"
+
+    def test_store_supports_writes(self, store: GpuMemoryStore) -> None:
+        assert store.supports_writes
+
+    def test_store_supports_listing(self, store: GpuMemoryStore) -> None:
+        assert store.supports_listing
+
+    def test_store_supports_partial_writes(self, store: GpuMemoryStore) -> None:
+        assert store.supports_partial_writes
+
+    def test_list_prefix(self, store: GpuMemoryStore) -> None:
+        assert True
+
+    def test_serizalizable_store(self, store: MemoryStore) -> None:
+        with pytest.raises(NotImplementedError):
+            store.__getstate__()
+
+        with pytest.raises(NotImplementedError):
+            store.__setstate__({})
+
+        with pytest.raises(NotImplementedError):
+            pickle.dumps(store)