Use dataclasses

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,7 +21,8 @@
 __all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
 
 
-class ExplicitRange(TypedDict):
+@dataclass
+class ExplicitRange:
     """Request a specific byte range"""
 
     start: int
@@ -29,21 +31,23 @@ class ExplicitRange(TypedDict):
     """The end of the byte range request (exclusive)."""
 
 
-class OffsetRange(TypedDict):
+@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 = None | ExplicitRange | OffsetRange | SuffixRange
+ByteRangeRequest: TypeAlias = ExplicitRange | OffsetRange | SuffixRange
 
 
 class Store(ABC):
@@ -190,7 +194,7 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         """Retrieve possibly partial values from given key_ranges.
 

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
@@ -505,7 +511,7 @@ async def _decode_partial_single(
                 if chunk_byte_slice:
                     chunk_bytes = await byte_getter.get(
                         prototype=chunk_spec.prototype,
-                        byte_range={"start": chunk_byte_slice[0], "end": chunk_byte_slice[1]},
+                        byte_range=ExplicitRange(chunk_byte_slice[0], chunk_byte_slice[1]),
                     )
                     if chunk_bytes:
                         shard_dict[chunk_coords] = chunk_bytes
@@ -697,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={"start": 0, "end": 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/_common.py

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

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:
@@ -199,7 +199,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: ByteRangeRequest = None,
+        byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
         if not self._is_open:
@@ -209,25 +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, dict):
-                if "start" in byte_range:
-                    value = prototype.buffer.from_bytes(
-                        await self.fs._cat_file(
-                            path,
-                            start=byte_range["start"],  # type: ignore[typeddict-item]
-                            end=byte_range["end"],  # type: ignore[typeddict-item]
-                        )
+            elif isinstance(byte_range, ExplicitRange):
+                value = prototype.buffer.from_bytes(
+                    await self.fs._cat_file(
+                        path,
+                        start=byte_range.start,
+                        end=byte_range.end,
                     )
-                elif "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)
-                    )
-                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:
@@ -276,7 +273,7 @@ async def exists(self, key: str) -> bool:
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         # docstring inherited
         if key_ranges:
@@ -290,18 +287,15 @@ async def get_partial_values(
                 if byte_range is None:
                     starts.append(None)
                     stops.append(None)
-                elif isinstance(byte_range, dict):
-                    if "start" in byte_range:
-                        starts.append(byte_range["start"])  # type: ignore[typeddict-item]
-                        stops.append(byte_range["end"])  # type: ignore[typeddict-item]
-                    elif "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
@@ -18,24 +18,21 @@
     from zarr.core.buffer import BufferPrototype
 
 
-def _get(path: Path, prototype: BufferPrototype, byte_range: ByteRangeRequest) -> Buffer:
+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 isinstance(byte_range, dict):
-            if "start" in byte_range:
-                f.seek(byte_range["start"])  # type: ignore[typeddict-item]
-                return prototype.buffer.from_bytes(f.read(byte_range["end"] - f.tell()))  # type: ignore[typeddict-item]
-            elif "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(
@@ -121,7 +118,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype | None = None,
-        byte_range: ByteRangeRequest = None,
+        byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
         if prototype is None:
@@ -139,7 +136,7 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         # docstring inherited
         args = []

src/zarr/storage/_logging.py

@@ -161,7 +161,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: ByteRangeRequest = None,
+        byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
         with self.log(key):
@@ -170,7 +170,7 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         # docstring inherited
         keys = ",".join([k[0] for k in key_ranges])

src/zarr/storage/_memory.py

@@ -75,7 +75,7 @@ async def get(
         self,
         key: str,
         prototype: BufferPrototype,
-        byte_range: ByteRangeRequest = None,
+        byte_range: ByteRangeRequest | None = None,
     ) -> Buffer | None:
         # docstring inherited
         if not self._is_open:
@@ -91,12 +91,12 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         # 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
@@ -45,21 +47,22 @@ def normalize_path(path: str | bytes | Path | None) -> str:
     return result
 
 
-def _normalize_byte_range_index(data: Buffer, byte_range: ByteRangeRequest) -> tuple[int, int]:
+def _normalize_byte_range_index(
+    data: Buffer, byte_range: ByteRangeRequest | None
+) -> tuple[int, int]:
     """
     Convert an ByteRangeRequest into an explicit start and stop
     """
     if byte_range is None:
         start = 0
         stop = len(data) + 1
-    elif "start" in byte_range:
-        # See https://github.com/python/mypy/issues/17087 for typeddict-item ignore explanation
-        start = byte_range["start"]  # type: ignore[typeddict-item]
-        stop = byte_range["end"]  # type: ignore[typeddict-item]
-    elif "offset" in byte_range:
-        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/_wrapper.py

@@ -77,7 +77,7 @@ async def get(
     async def get_partial_values(
         self,
         prototype: BufferPrototype,
-        key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+        key_ranges: Iterable[tuple[str, ByteRangeRequest | None]],
     ) -> list[Buffer | None]:
         return await self._store.get_partial_values(prototype, key_ranges)