Use TypedDicts for more literate ByteRangeRequests

Author: maxrjones

src/zarr/abc/store.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from asyncio import gather
 from itertools import starmap
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Protocol, TypedDict, runtime_checkable
 
 from zarr.core.buffer.core import default_buffer_prototype
 from zarr.core.common import concurrent_map
@@ -19,7 +19,22 @@
 
 __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
-ByteRangeRequest: TypeAlias = tuple[int | None, int | None]
+
+class OffsetRange(TypedDict):
+    """Request all bytes starting from a given byte offset"""
+
+    offset: int
+    """The byte offset for the offset range request."""
+
+
+class SuffixRange(TypedDict):
+    """Request up to the last `n` bytes"""
+
+    suffix: int
+    """The number of bytes from the suffix to request."""
+
+
+ByteRangeRequest: TypeAlias = None | tuple[int, int] | OffsetRange | SuffixRange
 
 
 class Store(ABC):
@@ -148,7 +163,13 @@ async def get(
         Parameters
         ----------
         key : str
-        byte_range : tuple[int | None, int | None], optional
+        byte_range : ByteRangeRequest, optional
+
+            The semantics of this argument are:
+
+            - tuple (int, int): Request a specific range of bytes (start, end). The end offset is exclusive. If the given range is zero-length or starts after the end of the object, an error will be returned. Additionally, if the range ends after the end of the object, the entire remainder of the object will be returned. Otherwise, the exact requested range will be returned.
+            - {"offset": int}: Request all bytes starting from a given byte offset. This is equivalent to bytes={int}- as an HTTP header.
+            - {"suffix": int}: Request the last int bytes. Note that here, int is the size of the request, not the byte offset. This is equivalent to bytes=-{int} as an HTTP header.
 
         Returns
         -------

src/zarr/storage/_utils.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from zarr.abc.store import ByteRangeRequest
     from zarr.core.buffer import Buffer
 
 
@@ -44,25 +45,19 @@ def normalize_path(path: str | bytes | Path | None) -> str:
     return result
 
 
-def _normalize_interval_index(
-    data: Buffer, interval: tuple[int | None, int | None] | None
-) -> tuple[int, int]:
+def _normalize_byte_range_index(data: Buffer, byte_range: ByteRangeRequest) -> tuple[int, int]:
     """
-    Convert an implicit interval into an explicit start and length
+    Convert an ByteRangeRequest into an explicit start and length
     """
-    if interval is None:
+    if byte_range is None:
         start = 0
         length = len(data)
-    else:
-        maybe_start, maybe_len = interval
-        if maybe_start is None:
-            start = 0
-        else:
-            start = maybe_start
-
-        if maybe_len is None:
-            length = len(data) - start
-        else:
-            length = maybe_len
-
+    elif isinstance(byte_range, tuple):
+        start = byte_range[0]
+        length = byte_range[1] - start
+    elif start := byte_range.get("offset"):
+        length = len(data) - start
+    elif suffix := byte_range.get("suffix"):
+        start = len(data) - suffix
+        length = len(data) - start
     return (start, length)

src/zarr/storage/common.py

@@ -102,7 +102,7 @@ async def open(
     async def get(
         self,
         prototype: BufferPrototype | None = None,
-        byte_range: ByteRangeRequest | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         """
         Read bytes from the store.

src/zarr/storage/fsspec.py

@@ -199,31 +199,33 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: ByteRangeRequest | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         # docstring inherited
         if not self._is_open:
             await self._open()
         path = _dereference_path(self.path, key)
 
         try:
-            if byte_range:
-                # fsspec uses start/end, not start/length
-                start, length = byte_range
-                if start is not None and length is not None:
-                    end = start + length
-                elif length is not None:
-                    end = length
-                else:
-                    end = None
-            value = prototype.buffer.from_bytes(
-                await (
-                    self.fs._cat_file(path, start=byte_range[0], end=end)
-                    if byte_range
-                    else self.fs._cat_file(path)
+            if byte_range is None:
+                value = prototype.buffer.from_bytes(await self.fs._cat_file(path))
+            elif isinstance(byte_range, tuple):
+                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)
+                    )
+                elif "suffix" in byte_range:
+                    value = prototype.buffer.from_bytes(
+                        await self.fs._cat_file(path, start=-byte_range["suffix"], end=None)
+                    )
+                else:
+                    raise ValueError("Invalid format for ByteRangeRequest")
+            else:
+                raise ValueError("Invalid format for ByteRangeRequest")
         except self.allowed_exceptions:
             return None
         except OSError as e:

src/zarr/storage/local.py

@@ -18,30 +18,25 @@
     from zarr.core.buffer import BufferPrototype
 
 
-def _get(
-    path: Path, prototype: BufferPrototype, byte_range: tuple[int | None, int | None] | None
-) -> Buffer:
-    if byte_range is not None:
-        if byte_range[0] is None:
-            start = 0
-        else:
-            start = byte_range[0]
-
-        end = (start + byte_range[1]) if byte_range[1] is not None else None
-    else:
+def _get(path: Path, prototype: BufferPrototype, byte_range: ByteRangeRequest | None) -> Buffer:
+    if byte_range is None:
         return prototype.buffer.from_bytes(path.read_bytes())
     with path.open("rb") as f:
         size = f.seek(0, io.SEEK_END)
-        if start is not None:
-            if start >= 0:
-                f.seek(start)
-            else:
-                f.seek(max(0, size + start))
-        if end is not None:
-            if end < 0:
-                end = size + end
+        if isinstance(byte_range, tuple):
+            start, end = byte_range
+            f.seek(start)
             return prototype.buffer.from_bytes(f.read(end - f.tell()))
-        return prototype.buffer.from_bytes(f.read())
+        elif isinstance(byte_range, dict):
+            if "offset" in byte_range:
+                f.seek(byte_range["offset"])
+            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())
+        else:
+            raise TypeError("Invalid format for ByteRangeRequest")
 
 
 def _put(
@@ -127,7 +122,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype | None = None,
-        byte_range: tuple[int | None, int | None] | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         # docstring inherited
         if prototype is None:

src/zarr/storage/logging.py

@@ -161,7 +161,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: tuple[int | None, int | None] | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         # docstring inherited
         with self.log(key):

src/zarr/storage/memory.py

@@ -6,7 +6,7 @@
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer, gpu
 from zarr.core.common import concurrent_map
-from zarr.storage._utils import _normalize_interval_index
+from zarr.storage._utils import _normalize_byte_range_index
 
 if TYPE_CHECKING:
     from collections.abc import AsyncIterator, Iterable, MutableMapping
@@ -75,15 +75,15 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: tuple[int | None, int | None] | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         # docstring inherited
         if not self._is_open:
             await self._open()
         assert isinstance(key, str)
         try:
             value = self._store_dict[key]
-            start, length = _normalize_interval_index(value, byte_range)
+            start, length = _normalize_byte_range_index(value, byte_range)
             return prototype.buffer.from_buffer(value[start : start + length])
         except KeyError:
             return None

src/zarr/storage/zip.py

@@ -138,23 +138,28 @@ def _get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: ByteRangeRequest | None = None,
+        byte_range: ByteRangeRequest = None,
     ) -> Buffer | None:
         # docstring inherited
         try:
             with self._zf.open(key) as f:  # will raise KeyError
                 if byte_range is None:
                     return prototype.buffer.from_bytes(f.read())
-                start, length = byte_range
-                if start:
-                    if start < 0:
-                        start = f.seek(start, os.SEEK_END) + start
+                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"])
+                    elif "suffix" in byte_range:
+                        f.seek(max(0, size - byte_range["suffix"]))
                     else:
-                        start = f.seek(start, os.SEEK_SET)
-                if length:
-                    return prototype.buffer.from_bytes(f.read(length))
-                else:
+                        raise TypeError("Invalid format for ByteRangeRequest")
                     return prototype.buffer.from_bytes(f.read())
+                else:
+                    raise TypeError("Invalid format for ByteRangeRequest")
         except KeyError:
             return None
 

src/zarr/testing/stateful.py

@@ -356,8 +356,7 @@ def get_partial_values(self, data: DataObject) -> None:
 
         for key, byte_range in key_range:
             start = byte_range[0] or 0
-            step = byte_range[1]
-            stop = start + step if step is not None else None
+            stop = byte_range[1]
             model_vals_ls.append(self.model[key][start:stop])
 
         assert all(

src/zarr/testing/store.py

@@ -17,7 +17,7 @@
 from zarr.abc.store import ByteRangeRequest, Store
 from zarr.core.buffer import Buffer, default_buffer_prototype
 from zarr.core.sync import _collect_aiterator
-from zarr.storage._utils import _normalize_interval_index
+from zarr.storage._utils import _normalize_byte_range_index
 from zarr.testing.utils import assert_bytes_equal
 
 __all__ = ["StoreTests"]
@@ -115,17 +115,15 @@ 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, (0, None), (1, None), (1, 2), (None, 1)])
-    async def test_get(
-        self, store: S, key: str, data: bytes, byte_range: tuple[int | None, int | None] | None
-    ) -> None:
+    @pytest.mark.parametrize("byte_range", [None, (1, 3), {"offset": 1}, {"suffix": 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.
         """
         data_buf = self.buffer_cls.from_bytes(data)
         await self.set(store, key, data_buf)
         observed = await store.get(key, prototype=default_buffer_prototype(), byte_range=byte_range)
-        start, length = _normalize_interval_index(data_buf, interval=byte_range)
+        start, length = _normalize_byte_range_index(data_buf, byte_range=byte_range)
         expected = data_buf[start : start + length]
         assert_bytes_equal(observed, expected)
 
@@ -180,12 +178,12 @@ async def test_set_many(self, store: S) -> None:
         [
             [],
             [("zarr.json", (0, 1))],
-            [("c/0", (0, 1)), ("zarr.json", (0, None))],
-            [("c/0/0", (0, 1)), ("c/0/1", (None, 2)), ("c/0/2", (0, 3))],
+            [("c/0", (0, 1)), ("zarr.json", None)],
+            [("c/0/0", (0, 1)), ("c/0/1", {"suffix": 2}), ("c/0/2", {"offset": 2})],
         ],
     )
     async def test_get_partial_values(
-        self, store: S, key_ranges: list[tuple[str, tuple[int | None, int | None]]]
+        self, store: S, key_ranges: list[tuple[str, ByteRangeRequest]]
     ) -> None:
         # put all of the data
         for key, _ in key_ranges: