Merge branch 'main' into perf/remove-isinstance-check

Author: d-v-b

changes/3657.bugfix.md

@@ -0,0 +1 @@
+Fix obstore _transform_list_dir implementation to correctly relativize paths (removing lstrip usage).

changes/3706.misc.md

@@ -0,0 +1 @@
+Allow NumPy ints as input when declaring a shape.

src/zarr/core/common.py

@@ -21,6 +21,7 @@
     overload,
 )
 
+import numpy as np
 from typing_extensions import ReadOnly
 
 from zarr.core.config import config as zarr_config
@@ -37,7 +38,7 @@
 ZMETADATA_V2_JSON = ".zmetadata"
 
 BytesLike = bytes | bytearray | memoryview
-ShapeLike = Iterable[int] | int
+ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any]
 # For backwards compatibility
 ChunkCoords = tuple[int, ...]
 ZarrFormat = Literal[2, 3]
@@ -185,23 +186,28 @@ def parse_named_configuration(
 
 
 def parse_shapelike(data: ShapeLike) -> tuple[int, ...]:
-    if isinstance(data, int):
+    """
+    Parse a shape-like input into an explicit shape.
+    """
+    if isinstance(data, int | np.integer):
         if data < 0:
             raise ValueError(f"Expected a non-negative integer. Got {data} instead")
-        return (data,)
+        return (int(data),)
     try:
         data_tuple = tuple(data)
     except TypeError as e:
         msg = f"Expected an integer or an iterable of integers. Got {data} instead."
         raise TypeError(msg) from e
 
-    if not all(isinstance(v, int) for v in data_tuple):
+    if not all(isinstance(v, int | np.integer) for v in data_tuple):
         msg = f"Expected an iterable of integers. Got {data} instead."
         raise TypeError(msg)
     if not all(v > -1 for v in data_tuple):
         msg = f"Expected all values to be non-negative. Got {data} instead."
         raise ValueError(msg)
-    return data_tuple
+
+    # cast NumPy scalars to plain python ints
+    return tuple(int(x) for x in data_tuple)
 
 
 def parse_fill_value(data: Any) -> Any:

src/zarr/storage/_obstore.py

@@ -4,6 +4,8 @@
 import contextlib
 import pickle
 from collections import defaultdict
+from itertools import chain
+from operator import itemgetter
 from typing import TYPE_CHECKING, Generic, Self, TypedDict, TypeVar
 
 from zarr.abc.store import (
@@ -15,6 +17,7 @@
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
+from zarr.storage._utils import _relativize_path
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence
@@ -263,10 +266,11 @@ async def _transform_list_dir(
     # We assume that the underlying object-store implementation correctly handles the
     # prefix, so we don't double-check that the returned results actually start with the
     # given prefix.
-    prefixes = [obj.lstrip(prefix).lstrip("/") for obj in list_result["common_prefixes"]]
-    objects = [obj["path"].removeprefix(prefix).lstrip("/") for obj in list_result["objects"]]
-    for item in prefixes + objects:
-        yield item
+    prefix = prefix.rstrip("/")
+    for path in chain(
+        list_result["common_prefixes"], map(itemgetter("path"), list_result["objects"])
+    ):
+        yield _relativize_path(path=path, prefix=prefix)
 
 
 class _BoundedRequest(TypedDict):

src/zarr/testing/store.py

@@ -492,24 +492,36 @@ async def test_list_empty_path(self, store: S) -> None:
         assert observed_prefix_sorted == expected_prefix_sorted
 
     async def test_list_dir(self, store: S) -> None:
-        root = "foo"
-        store_dict = {
-            root + "/zarr.json": self.buffer_cls.from_bytes(b"bar"),
-            root + "/c/1": self.buffer_cls.from_bytes(b"\x01"),
-        }
+        roots_and_keys: list[tuple[str, dict[str, Buffer]]] = [
+            (
+                "foo",
+                {
+                    "foo/zarr.json": self.buffer_cls.from_bytes(b"bar"),
+                    "foo/c/1": self.buffer_cls.from_bytes(b"\x01"),
+                },
+            ),
+            (
+                "foo/bar",
+                {
+                    "foo/bar/foobar_first_child": self.buffer_cls.from_bytes(b"1"),
+                    "foo/bar/foobar_second_child/zarr.json": self.buffer_cls.from_bytes(b"2"),
+                },
+            ),
+        ]
 
         assert await _collect_aiterator(store.list_dir("")) == ()
-        assert await _collect_aiterator(store.list_dir(root)) == ()
 
-        await store._set_many(store_dict.items())
+        for root, store_dict in roots_and_keys:
+            assert await _collect_aiterator(store.list_dir(root)) == ()
 
-        keys_observed = await _collect_aiterator(store.list_dir(root))
-        keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict}
+            await store._set_many(store_dict.items())
 
-        assert sorted(keys_observed) == sorted(keys_expected)
+            keys_observed = await _collect_aiterator(store.list_dir(root))
+            keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict}
+            assert sorted(keys_observed) == sorted(keys_expected)
 
-        keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
-        assert sorted(keys_expected) == sorted(keys_observed)
+            keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
+            assert sorted(keys_expected) == sorted(keys_observed)
 
     async def test_set_if_not_exists(self, store: S) -> None:
         key = "k"

tests/test_common.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from typing import TYPE_CHECKING, get_args
 
 import numpy as np
@@ -15,7 +16,6 @@
 from zarr.core.config import parse_indexing_order
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
     from typing import Any, Literal
 
 
@@ -115,9 +115,15 @@ def test_parse_shapelike_invalid_iterable_values(data: Any) -> None:
         parse_shapelike(data)
 
 
-@pytest.mark.parametrize("data", [range(10), [0, 1, 2, 3], (3, 4, 5), ()])
-def test_parse_shapelike_valid(data: Iterable[int]) -> None:
-    assert parse_shapelike(data) == tuple(data)
+@pytest.mark.parametrize(
+    "data", [range(10), [0, 1, 2, np.uint64(3)], (3, 4, 5), (), 1, np.uint8(1)]
+)
+def test_parse_shapelike_valid(data: Iterable[int] | int) -> None:
+    if isinstance(data, Iterable):
+        expected = tuple(data)
+    else:
+        expected = (data,)
+    assert parse_shapelike(data) == expected
 
 
 # todo: more dtypes