Merge branch 'byterange-dataclass' into literate-byte-ranges

Author: maxrjones

src/zarr/abc/store.py

@@ -2,8 +2,9 @@
 
 from abc import ABC, abstractmethod
 from asyncio import gather
+from dataclasses import dataclass
 from itertools import starmap
-from typing import TYPE_CHECKING, Protocol, TypedDict, runtime_checkable
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import concurrent_map
@@ -20,21 +21,33 @@
 __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
 
-class OffsetRange(TypedDict):
+@dataclass
+class ExplicitRange:
+    """Request a specific byte range"""
+
+    start: int
+    """The start of the byte range request (inclusive)."""
+    end: int
+    """The end of the byte range request (exclusive)."""
+
+
+@dataclass
+class OffsetRange:
     """Request all bytes starting from a given byte offset"""
 
     offset: int
     """The byte offset for the offset range request."""
 
 
-class SuffixRange(TypedDict):
+@dataclass
+class SuffixRange:
     """Request up to the last `n` bytes"""
 
     suffix: int
     """The number of bytes from the suffix to request."""
 
 
-ByteRangeRequest: TypeAlias = tuple[int, int] | OffsetRange | SuffixRange
+ByteRangeRequest: TypeAlias = ExplicitRange | OffsetRange | SuffixRange
 
 
 class Store(ABC):

src/zarr/codecs/sharding.py

@@ -17,7 +17,13 @@
     Codec,
     CodecPipeline,
 )
-from zarr.abc.store import ByteGetter, ByteRangeRequest, ByteSetter
+from zarr.abc.store import (
+    ByteGetter,
+    ByteRangeRequest,
+    ByteSetter,
+    ExplicitRange,
+    SuffixRange,
+)
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.core.array_spec import ArrayConfig, ArraySpec
@@ -504,7 +510,8 @@ async def _decode_partial_single(
                 chunk_byte_slice = shard_index.get_chunk_slice(chunk_coords)
                 if chunk_byte_slice:
                     chunk_bytes = await byte_getter.get(
-                        prototype=chunk_spec.prototype, byte_range=chunk_byte_slice
+                        prototype=chunk_spec.prototype,
+                        byte_range=ExplicitRange(chunk_byte_slice[0], chunk_byte_slice[1]),
                     )
                     if chunk_bytes:
                         shard_dict[chunk_coords] = chunk_bytes
@@ -696,11 +703,11 @@ async def _load_shard_index_maybe(
         shard_index_size = self._shard_index_size(chunks_per_shard)
         if self.index_location == ShardingCodecIndexLocation.start:
             index_bytes = await byte_getter.get(
-                prototype=numpy_buffer_prototype(), byte_range=(0, shard_index_size)
+                prototype=numpy_buffer_prototype(), byte_range=ExplicitRange(0, shard_index_size)
             )
         else:
             index_bytes = await byte_getter.get(
-                prototype=numpy_buffer_prototype(), byte_range={"suffix": shard_index_size}
+                prototype=numpy_buffer_prototype(), byte_range=SuffixRange(shard_index_size)
             )
         if index_bytes is not None:
             return await self._decode_shard_index(index_bytes, chunks_per_shard)

src/zarr/storage/_fsspec.py

@@ -3,7 +3,7 @@
 import warnings
 from typing import TYPE_CHECKING, Any
 
-from zarr.abc.store import ByteRangeRequest, Store
+from zarr.abc.store import ByteRangeRequest, ExplicitRange, OffsetRange, Store, SuffixRange
 from zarr.storage._common import _dereference_path
 
 if TYPE_CHECKING:
@@ -209,21 +209,22 @@ async def get(
         try:
             if byte_range is None:
                 value = prototype.buffer.from_bytes(await self.fs._cat_file(path))
-            elif isinstance(byte_range, tuple):
+            elif isinstance(byte_range, ExplicitRange):
                 value = prototype.buffer.from_bytes(
-                    await self.fs._cat_file(path, start=byte_range[0], end=byte_range[1])
-                )
-            elif isinstance(byte_range, dict):
-                if "offset" in byte_range:
-                    value = prototype.buffer.from_bytes(
-                        await self.fs._cat_file(path, start=byte_range["offset"], end=None)  # type: ignore[typeddict-item]
-                    )
-                elif "suffix" in byte_range:
-                    value = prototype.buffer.from_bytes(
-                        await self.fs._cat_file(path, start=-byte_range["suffix"], end=None)
+                    await self.fs._cat_file(
+                        path,
+                        start=byte_range.start,
+                        end=byte_range.end,
                     )
-                else:
-                    raise ValueError("Invalid format for ByteRangeRequest")
+                )
+            elif isinstance(byte_range, OffsetRange):
+                value = prototype.buffer.from_bytes(
+                    await self.fs._cat_file(path, start=byte_range.offset, end=None)
+                )
+            elif isinstance(byte_range, SuffixRange):
+                value = prototype.buffer.from_bytes(
+                    await self.fs._cat_file(path, start=-byte_range.suffix, end=None)
+                )
             else:
                 raise ValueError("Invalid format for ByteRangeRequest")
         except self.allowed_exceptions:
@@ -286,18 +287,15 @@ async def get_partial_values(
                 if byte_range is None:
                     starts.append(None)
                     stops.append(None)
-                elif isinstance(byte_range, tuple):
-                    starts.append(byte_range[0])
-                    stops.append(byte_range[1])
-                elif isinstance(byte_range, dict):
-                    if "offset" in byte_range:
-                        starts.append(byte_range["offset"])  # type: ignore[typeddict-item]
-                        stops.append(None)
-                    elif "suffix" in byte_range:
-                        starts.append(-byte_range["suffix"])
-                        stops.append(None)
-                    else:
-                        raise ValueError("Invalid format for ByteRangeRequest")
+                elif isinstance(byte_range, ExplicitRange):
+                    starts.append(byte_range.start)
+                    stops.append(byte_range.end)
+                elif isinstance(byte_range, OffsetRange):
+                    starts.append(byte_range.offset)
+                    stops.append(None)
+                elif isinstance(byte_range, SuffixRange):
+                    starts.append(-byte_range.suffix)
+                    stops.append(None)
                 else:
                     raise ValueError("Invalid format for ByteRangeRequest")
         else:

src/zarr/storage/_local.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from zarr.abc.store import ByteRangeRequest, Store
+from zarr.abc.store import ByteRangeRequest, ExplicitRange, OffsetRange, Store, SuffixRange
 from zarr.core.buffer import Buffer
 from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import concurrent_map
@@ -23,20 +23,16 @@ def _get(path: Path, prototype: BufferPrototype, byte_range: ByteRangeRequest |
         return prototype.buffer.from_bytes(path.read_bytes())
     with path.open("rb") as f:
         size = f.seek(0, io.SEEK_END)
-        if isinstance(byte_range, tuple):
-            start, end = byte_range
-            f.seek(start)
-            return prototype.buffer.from_bytes(f.read(end - f.tell()))
-        elif isinstance(byte_range, dict):
-            if "offset" in byte_range:
-                f.seek(byte_range["offset"])  # type: ignore[typeddict-item]
-            elif "suffix" in byte_range:
-                f.seek(max(0, size - byte_range["suffix"]))
-            else:
-                raise TypeError("Invalid format for ByteRangeRequest")
-            return prototype.buffer.from_bytes(f.read())
+        if isinstance(byte_range, ExplicitRange):
+            f.seek(byte_range.start)
+            return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell()))
+        elif isinstance(byte_range, OffsetRange):
+            f.seek(byte_range.offset)
+        elif isinstance(byte_range, SuffixRange):
+            f.seek(max(0, size - byte_range.suffix))
         else:
             raise TypeError("Invalid format for ByteRangeRequest")
+        return prototype.buffer.from_bytes(f.read())
 
 
 def _put(

src/zarr/storage/_memory.py

@@ -96,7 +96,7 @@ async def get_partial_values(
         # docstring inherited
 
         # All the key-ranges arguments goes with the same prototype
-        async def _get(key: str, byte_range: ByteRangeRequest) -> Buffer | None:
+        async def _get(key: str, byte_range: ByteRangeRequest | None) -> Buffer | None:
             return await self.get(key, prototype=prototype, byte_range=byte_range)
 
         return await concurrent_map(key_ranges, _get, limit=None)

src/zarr/storage/_utils.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+from zarr.abc.store import ExplicitRange, OffsetRange, SuffixRange
+
 if TYPE_CHECKING:
     from zarr.abc.store import ByteRangeRequest
     from zarr.core.buffer import Buffer
@@ -54,14 +56,13 @@ def _normalize_byte_range_index(
     if byte_range is None:
         start = 0
         stop = len(data) + 1
-    elif isinstance(byte_range, tuple):
-        start = byte_range[0]
-        stop = byte_range[1]
-    elif "offset" in byte_range:
-        # See https://github.com/python/mypy/issues/17087 for typeddict-item ignore explanation
-        start = byte_range["offset"]  # type: ignore[typeddict-item]
+    elif isinstance(byte_range, ExplicitRange):
+        start = byte_range.start
+        stop = byte_range.end
+    elif isinstance(byte_range, OffsetRange):
+        start = byte_range.offset
         stop = len(data) + 1
-    elif "suffix" in byte_range:
-        start = len(data) - byte_range["suffix"]
+    elif isinstance(byte_range, SuffixRange):
+        start = len(data) - byte_range.suffix
         stop = len(data) + 1
     return (start, stop)

src/zarr/storage/_zip.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
-from zarr.abc.store import ByteRangeRequest, Store
+from zarr.abc.store import ByteRangeRequest, ExplicitRange, OffsetRange, Store, SuffixRange
 from zarr.core.buffer import Buffer, BufferPrototype
 
 if TYPE_CHECKING:
@@ -145,21 +145,17 @@ def _get(
             with self._zf.open(key) as f:  # will raise KeyError
                 if byte_range is None:
                     return prototype.buffer.from_bytes(f.read())
-                if isinstance(byte_range, tuple):
-                    start, end = byte_range
-                    f.seek(start)
-                    return prototype.buffer.from_bytes(f.read(end - f.tell()))
-                elif isinstance(byte_range, dict):
-                    size = f.seek(0, os.SEEK_END)
-                    if "offset" in byte_range:
-                        f.seek(byte_range["offset"])  # type: ignore[typeddict-item]
-                    elif "suffix" in byte_range:
-                        f.seek(max(0, size - byte_range["suffix"]))
-                    else:
-                        raise TypeError("Invalid format for ByteRangeRequest")
-                    return prototype.buffer.from_bytes(f.read())
+                elif isinstance(byte_range, ExplicitRange):
+                    f.seek(byte_range.start)
+                    return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell()))
+                size = f.seek(0, os.SEEK_END)
+                if isinstance(byte_range, OffsetRange):
+                    f.seek(byte_range.offset)
+                elif isinstance(byte_range, SuffixRange):
+                    f.seek(max(0, size - byte_range.suffix))
                 else:
                     raise TypeError("Invalid format for ByteRangeRequest")
+                return prototype.buffer.from_bytes(f.read())
         except KeyError:
             return None
 

src/zarr/testing/stateful.py

@@ -355,8 +355,8 @@ def get_partial_values(self, data: DataObject) -> None:
         model_vals_ls = []
 
         for key, byte_range in key_range:
-            start = byte_range[0] or 0
-            stop = byte_range[1]
+            start = byte_range.start
+            stop = byte_range.end
             model_vals_ls.append(self.model[key][start:stop])
 
         assert all(

src/zarr/testing/store.py

@@ -14,7 +14,7 @@
 
 import pytest
 
-from zarr.abc.store import ByteRangeRequest, Store
+from zarr.abc.store import ByteRangeRequest, ExplicitRange, OffsetRange, Store, SuffixRange
 from zarr.core.buffer import Buffer, default_buffer_prototype
 from zarr.core.sync import _collect_aiterator
 from zarr.storage._utils import _normalize_byte_range_index
@@ -115,7 +115,9 @@ def test_store_supports_listing(self, store: S) -> None:
 
     @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"])
     @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""])
-    @pytest.mark.parametrize("byte_range", [None, (1, 3), {"offset": 1}, {"suffix": 1}])
+    @pytest.mark.parametrize(
+        "byte_range", [None, ExplicitRange(1, 4), OffsetRange(1), SuffixRange(1)]
+    )
     async def test_get(self, store: S, key: str, data: bytes, byte_range: ByteRangeRequest) -> None:
         """
         Ensure that data can be read from the store using the store.get method.
@@ -177,9 +179,9 @@ async def test_set_many(self, store: S) -> None:
         "key_ranges",
         [
             [],
-            [("zarr.json", (0, 1))],
-            [("c/0", (0, 1)), ("zarr.json", None)],
-            [("c/0/0", (0, 1)), ("c/0/1", {"suffix": 2}), ("c/0/2", {"offset": 2})],
+            [("zarr.json", ExplicitRange(0, 2))],
+            [("c/0", ExplicitRange(0, 2)), ("zarr.json", None)],
+            [("c/0/0", ExplicitRange(0, 2)), ("c/0/1", SuffixRange(2)), ("c/0/2", OffsetRange(2))],
         ],
     )
     async def test_get_partial_values(

src/zarr/testing/strategies.py

@@ -7,6 +7,7 @@
 from hypothesis.strategies import SearchStrategy
 
 import zarr
+from zarr.abc.store import ExplicitRange
 from zarr.core.array import Array
 from zarr.core.common import ZarrFormat
 from zarr.core.sync import sync
@@ -197,9 +198,10 @@ def key_ranges(
         [(key, (range_start, range_end)),
          (key, (range_start, range_end)),...]
     """
-    byte_ranges = st.tuples(
-        st.integers(min_value=0, max_value=max_size),
-        st.integers(min_value=0, max_value=max_size),
+    byte_ranges = st.builds(
+        ExplicitRange,
+        start=st.integers(min_value=0, max_value=max_size),
+        end=st.integers(min_value=0, max_value=max_size),
     )
     key_tuple = st.tuples(keys, byte_ranges)
     return st.lists(key_tuple, min_size=1, max_size=10)