From 021bd440c3ee2fae13aaa255d722bb582d22b6a0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 4 Jan 2026 17:29:56 +0100 Subject: [PATCH 1/7] add store routines for getting bytes and json --- src/zarr/abc/store.py | 215 ++++++++++++++++++++++++++++- src/zarr/storage/_common.py | 107 +++++++++++++++ src/zarr/storage/_local.py | 232 +++++++++++++++++++++++++++++++- src/zarr/storage/_memory.py | 232 +++++++++++++++++++++++++++++++- src/zarr/testing/store.py | 41 +++++- tests/test_store/test_local.py | 82 +++++++++++ tests/test_store/test_memory.py | 70 ++++++++++ 7 files changed, 974 insertions(+), 5 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index 4b3edf78d1..7d0589c836 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -1,11 +1,14 @@ from __future__ import annotations +import asyncio +import json from abc import ABC, abstractmethod -from asyncio import gather from dataclasses import dataclass from itertools import starmap from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable +from zarr.core.sync import sync + if TYPE_CHECKING: from collections.abc import AsyncGenerator, AsyncIterator, Iterable from types import TracebackType @@ -206,6 +209,214 @@ async def get( """ ... + async def get_bytes_async( + self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> bytes: + """ + Retrieve raw bytes from the store asynchronously. + + This is a convenience method that wraps ``get()`` and converts the result + to bytes. Use this when you need the raw byte content of a stored value. + + Parameters + ---------- + key : str + The key identifying the data to retrieve. + prototype : BufferPrototype + The buffer prototype to use for reading the data. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + See Also + -------- + get : Lower-level method that returns a Buffer object. + get_bytes : Synchronous version of this method. + get_json_async : Asynchronous method for retrieving and parsing JSON data. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> await store.set("data", Buffer.from_bytes(b"hello world")) + >>> data = await store.get_bytes_async("data", prototype=default_buffer_prototype()) + >>> print(data) + b'hello world' + """ + buffer = await self.get(key, prototype, byte_range) + if buffer is None: + raise FileNotFoundError(key) + return buffer.to_bytes() + + def get_bytes( + self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> bytes: + """ + Retrieve raw bytes from the store synchronously. + + This is a synchronous wrapper around ``get_bytes_async()``. It should only + be called from non-async code. For async contexts, use ``get_bytes_async()`` + instead. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype + The buffer prototype to use for reading the data. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + Warnings + -------- + Do not call this method from async functions. Use ``get_bytes_async()`` instead + to avoid blocking the event loop. + + See Also + -------- + get_bytes_async : Asynchronous version of this method. + get_json : Synchronous method for retrieving and parsing JSON data. + + Examples + -------- + >>> store = MemoryStore() + >>> store.set("data", Buffer.from_bytes(b"hello world")) + >>> data = store.get_bytes("data", prototype=default_buffer_prototype()) + >>> print(data) + b'hello world' + """ + + return sync(self.get_bytes_async(key, prototype=prototype, byte_range=byte_range)) + + async def get_json_async( + self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Any: + """ + Retrieve and parse JSON data from the store asynchronously. + + This is a convenience method that retrieves bytes from the store and + parses them as JSON. Commonly used for reading Zarr metadata files + like ``zarr.json``. + + Parameters + ---------- + key : str + The key identifying the JSON data to retrieve. + prototype : BufferPrototype + The buffer prototype to use for reading the data. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + get_bytes_async : Method for retrieving raw bytes without parsing. + get_json : Synchronous version of this method. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> data = await store.get_json_async("zarr.json", prototype=default_buffer_prototype()) + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + + return json.loads( + await self.get_bytes_async(key, prototype=prototype, byte_range=byte_range) + ) + + def get_json( + self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Any: + """ + Retrieve and parse JSON data from the store synchronously. + + This is a synchronous wrapper around ``get_json_async()``. It should only + be called from non-async code. For async contexts, use ``get_json_async()`` + instead. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype + The buffer prototype to use for reading the data. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + Warnings + -------- + Do not call this method from async functions. Use ``get_json_async()`` instead + to avoid blocking the event loop. + + See Also + -------- + get_json_async : Asynchronous version of this method. + get_bytes : Synchronous method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = MemoryStore() + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> data = store.get_json("zarr.json", prototype=default_buffer_prototype()) + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + + return sync(self.get_json_async(key, prototype=prototype, byte_range=byte_range)) + @abstractmethod async def get_partial_values( self, @@ -278,7 +489,7 @@ async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None: """ Insert multiple (key, value) pairs into storage. """ - await gather(*starmap(self.set, values)) + await asyncio.gather(*starmap(self.set, values)) @property def supports_consolidated_metadata(self) -> bool: diff --git a/src/zarr/storage/_common.py b/src/zarr/storage/_common.py index d762097cc3..b8dcaff3be 100644 --- a/src/zarr/storage/_common.py +++ b/src/zarr/storage/_common.py @@ -228,6 +228,113 @@ async def is_empty(self) -> bool: """ return await self.store.is_empty(self.path) + async def get_bytes_async( + self, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the store path asynchronously. + + This is a convenience method that wraps ``get()`` and converts the result + to bytes. The ``prototype`` parameter is optional and defaults to the + standard buffer prototype. + + Parameters + ---------- + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + + Returns + ------- + bytes + The raw bytes stored at this path. + + Raises + ------ + FileNotFoundError + If the path does not exist in the store. + + See Also + -------- + get : Lower-level method that returns a Buffer object. + get_json_async : Asynchronous method for retrieving and parsing JSON data. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> path = StorePath(store, "data") + >>> await path.set(Buffer.from_bytes(b"hello world")) + >>> data = await path.get_bytes_async() + >>> print(data) + b'hello world' + """ + if prototype is None: + prototype = default_buffer_prototype() + return await self.store.get_bytes_async( + self.path, prototype=prototype, byte_range=byte_range + ) + + async def get_json_async( + self, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the store path asynchronously. + + This is a convenience method that retrieves bytes from the store and + parses them as JSON. The ``prototype`` parameter is optional and defaults + to the standard buffer prototype. + + Parameters + ---------- + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the path does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + get_bytes_async : Method for retrieving raw bytes without parsing. + get : Lower-level method that returns a Buffer object. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> path = StorePath(store, "zarr.json") + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await path.set(Buffer.from_bytes(json.dumps(metadata).encode())) + >>> data = await path.get_json_async() + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return await self.store.get_json_async( + self.path, prototype=prototype, byte_range=byte_range + ) + def __truediv__(self, other: str) -> StorePath: """Combine this store path with another path""" return self.__class__(self.store, _dereference_path(self.path, other)) diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index f64da71bb4..13c86a2f22 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -8,7 +8,7 @@ import sys import uuid from pathlib import Path -from typing import TYPE_CHECKING, BinaryIO, Literal, Self +from typing import TYPE_CHECKING, Any, BinaryIO, Literal, Self from zarr.abc.store import ( ByteRequest, @@ -306,6 +306,236 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: except (FileNotFoundError, NotADirectoryError): pass + async def get_bytes_async( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the local store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes_async`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + See Also + -------- + Store.get_bytes_async : Base implementation with full documentation. + get_bytes : Synchronous version of this method. + + Examples + -------- + >>> store = await LocalStore.open("data") + >>> await store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for LocalStore + >>> data = await store.get_bytes_async("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super().get_bytes_async(key, prototype=prototype, byte_range=byte_range) + + def get_bytes( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the local store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + Warnings + -------- + Do not call this method from async functions. Use ``get_bytes_async()`` instead. + + See Also + -------- + Store.get_bytes : Base implementation with full documentation. + get_bytes_async : Asynchronous version of this method. + + Examples + -------- + >>> store = LocalStore("data") + >>> store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for LocalStore + >>> data = store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return super().get_bytes(key, prototype=prototype, byte_range=byte_range) + + async def get_json_async( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the local store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json_async`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + Store.get_json_async : Base implementation with full documentation. + get_json : Synchronous version of this method. + get_bytes_async : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = await LocalStore.open("data") + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for LocalStore + >>> data = await store.get_json_async("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super().get_json_async(key, prototype=prototype, byte_range=byte_range) + + def get_json( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the local store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + Warnings + -------- + Do not call this method from async functions. Use ``get_json_async()`` instead. + + See Also + -------- + Store.get_json : Base implementation with full documentation. + get_json_async : Asynchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = LocalStore("data") + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for LocalStore + >>> data = store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return super().get_json(key, prototype=prototype, byte_range=byte_range) + async def move(self, dest_root: Path | str) -> None: """ Move the store to another path. The old root directory is deleted. diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index 904be922d7..b56771f62a 100644 --- a/src/zarr/storage/_memory.py +++ b/src/zarr/storage/_memory.py @@ -1,7 +1,7 @@ from __future__ import annotations from logging import getLogger -from typing import TYPE_CHECKING, Self +from typing import TYPE_CHECKING, Any, Self from zarr.abc.store import ByteRequest, Store from zarr.core.buffer import Buffer, gpu @@ -175,6 +175,236 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: for key in keys_unique: yield key + async def get_bytes_async( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the memory store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes_async`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + See Also + -------- + Store.get_bytes_async : Base implementation with full documentation. + get_bytes : Synchronous version of this method. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> await store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for MemoryStore + >>> data = await store.get_bytes_async("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super().get_bytes_async(key, prototype=prototype, byte_range=byte_range) + + def get_bytes( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> bytes: + """ + Retrieve raw bytes from the memory store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + + Returns + ------- + bytes + The raw bytes stored at the given key. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + + Warnings + -------- + Do not call this method from async functions. Use ``get_bytes_async()`` instead. + + See Also + -------- + Store.get_bytes : Base implementation with full documentation. + get_bytes_async : Asynchronous version of this method. + + Examples + -------- + >>> store = MemoryStore() + >>> store.set("data", Buffer.from_bytes(b"hello")) + >>> # No need to specify prototype for MemoryStore + >>> data = store.get_bytes("data") + >>> print(data) + b'hello' + """ + if prototype is None: + prototype = default_buffer_prototype() + return super().get_bytes(key, prototype=prototype, byte_range=byte_range) + + async def get_json_async( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the memory store asynchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json_async`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + See Also + -------- + Store.get_json_async : Base implementation with full documentation. + get_json : Synchronous version of this method. + get_bytes_async : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = await MemoryStore.open() + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for MemoryStore + >>> data = await store.get_json_async("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return await super().get_json_async(key, prototype=prototype, byte_range=byte_range) + + def get_json( + self, + key: str = "", + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Any: + """ + Retrieve and parse JSON data from the memory store synchronously. + + This is a convenience override that makes the ``prototype`` parameter optional + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` + for full documentation. + + Parameters + ---------- + key : str, optional + The key identifying the JSON data to retrieve. Defaults to an empty string. + prototype : BufferPrototype, optional + The buffer prototype to use for reading the data. If None, uses + ``default_buffer_prototype()``. + byte_range : ByteRequest, optional + If specified, only retrieve a portion of the stored data. + Note: Using byte ranges with JSON may result in invalid JSON. + + Returns + ------- + Any + The parsed JSON data. This follows the behavior of ``json.loads()`` and + can be any JSON-serializable type: dict, list, str, int, float, bool, or None. + + Raises + ------ + FileNotFoundError + If the key does not exist in the store. + json.JSONDecodeError + If the stored data is not valid JSON. + + Warnings + -------- + Do not call this method from async functions. Use ``get_json_async()`` instead. + + See Also + -------- + Store.get_json : Base implementation with full documentation. + get_json_async : Asynchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. + + Examples + -------- + >>> store = MemoryStore() + >>> import json + >>> metadata = {"zarr_format": 3, "node_type": "array"} + >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) + >>> # No need to specify prototype for MemoryStore + >>> data = store.get_json("zarr.json") + >>> print(data) + {'zarr_format': 3, 'node_type': 'array'} + """ + if prototype is None: + prototype = default_buffer_prototype() + return super().get_json(key, prototype=prototype, byte_range=byte_range) + class GpuMemoryStore(MemoryStore): """ diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index ad3b80da41..bee28639a2 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import json import pickle from abc import abstractmethod from typing import TYPE_CHECKING, Generic, TypeVar @@ -23,7 +24,7 @@ SuffixByteRequest, ) from zarr.core.buffer import Buffer, default_buffer_prototype -from zarr.core.sync import _collect_aiterator +from zarr.core.sync import _collect_aiterator, sync from zarr.storage._utils import _normalize_byte_range_index from zarr.testing.utils import assert_bytes_equal @@ -526,6 +527,44 @@ async def test_set_if_not_exists(self, store: S) -> None: result = await store.get("k2", default_buffer_prototype()) assert result == new + async def test_get_bytes_async(self, store: S) -> None: + """ + Test that the get_bytes_async method reads bytes. + """ + data = b"hello world" + key = "zarr.json" + await self.set(store, key, self.buffer_cls.from_bytes(data)) + assert await store.get_bytes_async(key, prototype=default_buffer_prototype()) == data + + def test_get_bytes_sync(self, store: S) -> None: + """ + Test that the get_bytes method reads bytes. + """ + data = b"hello world" + key = "zarr.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(data))) + assert store.get_bytes(key, prototype=default_buffer_prototype()) == data + + async def test_get_json_async(self, store: S) -> None: + """ + Test that the get_bytes_async method reads json. + """ + data = {"foo": "bar"} + data_bytes = json.dumps(data).encode("utf-8") + key = "zarr.json" + await self.set(store, key, self.buffer_cls.from_bytes(data_bytes)) + assert await store.get_json_async(key, prototype=default_buffer_prototype()) == data + + def test_get_json_sync(self, store: S) -> None: + """ + Test that the get_json method reads json. + """ + data = {"foo": "bar"} + data_bytes = json.dumps(data).encode("utf-8") + key = "zarr.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(data_bytes))) + assert store.get_json(key, prototype=default_buffer_prototype()) == data + class LatencyStore(WrapperStore[Store]): """ diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index 6756bc83d9..35d48e3f95 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -150,3 +150,85 @@ def test_atomic_write_exclusive_preexisting(tmp_path: pathlib.Path) -> None: f.write(b"abc") assert path.read_bytes() == b"xyz" assert list(path.parent.iterdir()) == [path] # no temp files + + +async def test_get_bytes_with_prototype_none(tmp_path: pathlib.Path) -> None: + """Test that get_bytes_async works with prototype=None.""" + from zarr.core.buffer import cpu + from zarr.core.buffer.core import default_buffer_prototype + + store = await LocalStore.open(root=tmp_path) + data = b"hello world" + key = "test_key" + await store.set(key, cpu.Buffer.from_bytes(data)) + + # Test with None (default) + result_none = await store.get_bytes_async(key) + assert result_none == data + + # Test with explicit prototype + result_proto = await store.get_bytes_async(key, prototype=default_buffer_prototype()) + assert result_proto == data + + +def test_get_bytes_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: + """Test that get_bytes works with prototype=None.""" + from zarr.core.buffer import cpu + from zarr.core.buffer.core import default_buffer_prototype + from zarr.core.sync import sync + + store = sync(LocalStore.open(root=tmp_path)) + data = b"hello world" + key = "test_key" + sync(store.set(key, cpu.Buffer.from_bytes(data))) + + # Test with None (default) + result_none = store.get_bytes(key) + assert result_none == data + + # Test with explicit prototype + result_proto = store.get_bytes(key, prototype=default_buffer_prototype()) + assert result_proto == data + + +async def test_get_json_with_prototype_none(tmp_path: pathlib.Path) -> None: + """Test that get_json_async works with prototype=None.""" + import json + + from zarr.core.buffer import cpu + from zarr.core.buffer.core import default_buffer_prototype + + store = await LocalStore.open(root=tmp_path) + data = {"foo": "bar", "number": 42} + key = "test.json" + await store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode())) + + # Test with None (default) + result_none = await store.get_json_async(key) + assert result_none == data + + # Test with explicit prototype + result_proto = await store.get_json_async(key, prototype=default_buffer_prototype()) + assert result_proto == data + + +def test_get_json_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: + """Test that get_json works with prototype=None.""" + import json + + from zarr.core.buffer import cpu + from zarr.core.buffer.core import default_buffer_prototype + from zarr.core.sync import sync + + store = sync(LocalStore.open(root=tmp_path)) + data = {"foo": "bar", "number": 42} + key = "test.json" + sync(store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode()))) + + # Test with None (default) + result_none = store.get_json(key) + assert result_none == data + + # Test with explicit prototype + result_proto = store.get_json(key, prototype=default_buffer_prototype()) + assert result_proto == data diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index 29fa9b2964..b56d9933d4 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -76,6 +76,76 @@ async def test_deterministic_size( np.testing.assert_array_equal(a[:3], 1) np.testing.assert_array_equal(a[3:], 0) + async def test_get_bytes_with_prototype_none(self, store: MemoryStore) -> None: + """Test that get_bytes_async works with prototype=None.""" + from zarr.core.buffer.core import default_buffer_prototype + + data = b"hello world" + key = "test_key" + await self.set(store, key, self.buffer_cls.from_bytes(data)) + + # Test with None (default) + result_none = await store.get_bytes_async(key) + assert result_none == data + + # Test with explicit prototype + result_proto = await store.get_bytes_async(key, prototype=default_buffer_prototype()) + assert result_proto == data + + def test_get_bytes_sync_with_prototype_none(self, store: MemoryStore) -> None: + """Test that get_bytes works with prototype=None.""" + from zarr.core.buffer.core import default_buffer_prototype + from zarr.core.sync import sync + + data = b"hello world" + key = "test_key" + sync(self.set(store, key, self.buffer_cls.from_bytes(data))) + + # Test with None (default) + result_none = store.get_bytes(key) + assert result_none == data + + # Test with explicit prototype + result_proto = store.get_bytes(key, prototype=default_buffer_prototype()) + assert result_proto == data + + async def test_get_json_with_prototype_none(self, store: MemoryStore) -> None: + """Test that get_json_async works with prototype=None.""" + import json + + from zarr.core.buffer.core import default_buffer_prototype + + data = {"foo": "bar", "number": 42} + key = "test.json" + await self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode())) + + # Test with None (default) + result_none = await store.get_json_async(key) + assert result_none == data + + # Test with explicit prototype + result_proto = await store.get_json_async(key, prototype=default_buffer_prototype()) + assert result_proto == data + + def test_get_json_sync_with_prototype_none(self, store: MemoryStore) -> None: + """Test that get_json works with prototype=None.""" + import json + + from zarr.core.buffer.core import default_buffer_prototype + from zarr.core.sync import sync + + data = {"foo": "bar", "number": 42} + key = "test.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode()))) + + # Test with None (default) + result_none = store.get_json(key) + assert result_none == data + + # Test with explicit prototype + result_proto = store.get_json(key, prototype=default_buffer_prototype()) + assert result_proto == data + # TODO: fix this warning @pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") From 7d26b8ee4d33c6c784e42784e1acf37c5836dd8e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 4 Jan 2026 18:02:29 +0100 Subject: [PATCH 2/7] check for FileNotFoundError when a key is missing --- src/zarr/testing/store.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index bee28639a2..30ff376fb0 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -535,6 +535,8 @@ async def test_get_bytes_async(self, store: S) -> None: key = "zarr.json" await self.set(store, key, self.buffer_cls.from_bytes(data)) assert await store.get_bytes_async(key, prototype=default_buffer_prototype()) == data + with pytest.raises(FileNotFoundError): + await store.get_bytes_async("nonexistent_key", prototype=default_buffer_prototype()) def test_get_bytes_sync(self, store: S) -> None: """ From 971c3e4fb6dd6c895e49c2763be5c1b9164b9114 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 4 Jan 2026 18:02:39 +0100 Subject: [PATCH 3/7] remove storepath methods --- src/zarr/storage/_common.py | 107 ------------------------------------ 1 file changed, 107 deletions(-) diff --git a/src/zarr/storage/_common.py b/src/zarr/storage/_common.py index b8dcaff3be..d762097cc3 100644 --- a/src/zarr/storage/_common.py +++ b/src/zarr/storage/_common.py @@ -228,113 +228,6 @@ async def is_empty(self) -> bool: """ return await self.store.is_empty(self.path) - async def get_bytes_async( - self, - prototype: BufferPrototype | None = None, - byte_range: ByteRequest | None = None, - ) -> bytes: - """ - Retrieve raw bytes from the store path asynchronously. - - This is a convenience method that wraps ``get()`` and converts the result - to bytes. The ``prototype`` parameter is optional and defaults to the - standard buffer prototype. - - Parameters - ---------- - prototype : BufferPrototype, optional - The buffer prototype to use for reading the data. If None, uses - ``default_buffer_prototype()``. - byte_range : ByteRequest, optional - If specified, only retrieve a portion of the stored data. - Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. - - Returns - ------- - bytes - The raw bytes stored at this path. - - Raises - ------ - FileNotFoundError - If the path does not exist in the store. - - See Also - -------- - get : Lower-level method that returns a Buffer object. - get_json_async : Asynchronous method for retrieving and parsing JSON data. - - Examples - -------- - >>> store = await MemoryStore.open() - >>> path = StorePath(store, "data") - >>> await path.set(Buffer.from_bytes(b"hello world")) - >>> data = await path.get_bytes_async() - >>> print(data) - b'hello world' - """ - if prototype is None: - prototype = default_buffer_prototype() - return await self.store.get_bytes_async( - self.path, prototype=prototype, byte_range=byte_range - ) - - async def get_json_async( - self, - prototype: BufferPrototype | None = None, - byte_range: ByteRequest | None = None, - ) -> Any: - """ - Retrieve and parse JSON data from the store path asynchronously. - - This is a convenience method that retrieves bytes from the store and - parses them as JSON. The ``prototype`` parameter is optional and defaults - to the standard buffer prototype. - - Parameters - ---------- - prototype : BufferPrototype, optional - The buffer prototype to use for reading the data. If None, uses - ``default_buffer_prototype()``. - byte_range : ByteRequest, optional - If specified, only retrieve a portion of the stored data. - Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``. - Note: Using byte ranges with JSON may result in invalid JSON. - - Returns - ------- - Any - The parsed JSON data. This follows the behavior of ``json.loads()`` and - can be any JSON-serializable type: dict, list, str, int, float, bool, or None. - - Raises - ------ - FileNotFoundError - If the path does not exist in the store. - json.JSONDecodeError - If the stored data is not valid JSON. - - See Also - -------- - get_bytes_async : Method for retrieving raw bytes without parsing. - get : Lower-level method that returns a Buffer object. - - Examples - -------- - >>> store = await MemoryStore.open() - >>> path = StorePath(store, "zarr.json") - >>> metadata = {"zarr_format": 3, "node_type": "array"} - >>> await path.set(Buffer.from_bytes(json.dumps(metadata).encode())) - >>> data = await path.get_json_async() - >>> print(data) - {'zarr_format': 3, 'node_type': 'array'} - """ - if prototype is None: - prototype = default_buffer_prototype() - return await self.store.get_json_async( - self.path, prototype=prototype, byte_range=byte_range - ) - def __truediv__(self, other: str) -> StorePath: """Combine this store path with another path""" return self.__class__(self.store, _dereference_path(self.path, other)) From d70a5e5277ca7a225d9fb0ac22945b5e1fcb8cc3 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 4 Jan 2026 18:24:39 +0100 Subject: [PATCH 4/7] changelog --- changes/3638.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/3638.feature.md diff --git a/changes/3638.feature.md b/changes/3638.feature.md new file mode 100644 index 0000000000..ad2276fd51 --- /dev/null +++ b/changes/3638.feature.md @@ -0,0 +1 @@ +Add methods for reading stored objects as bytes and JSON-decoded bytes to store classes. \ No newline at end of file From a21305887b7287d957811afb0a8ec8f894d8cfc7 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 8 Jan 2026 10:37:23 +0100 Subject: [PATCH 5/7] rename methods --- src/zarr/abc/store.py | 45 +++++++++++++------------- src/zarr/storage/_local.py | 26 +++++++-------- src/zarr/storage/_memory.py | 26 +++++++-------- src/zarr/testing/store.py | 12 +++---- tests/test_store/test_local.py | 56 ++++++++++++++------------------- tests/test_store/test_memory.py | 18 +++++------ 6 files changed, 85 insertions(+), 98 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index 7d0589c836..e685e4b3b0 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -209,7 +209,7 @@ async def get( """ ... - async def get_bytes_async( + async def get_bytes( self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None ) -> bytes: """ @@ -242,7 +242,7 @@ async def get_bytes_async( -------- get : Lower-level method that returns a Buffer object. get_bytes : Synchronous version of this method. - get_json_async : Asynchronous method for retrieving and parsing JSON data. + get_json : Asynchronous method for retrieving and parsing JSON data. Examples -------- @@ -257,7 +257,7 @@ async def get_bytes_async( raise FileNotFoundError(key) return buffer.to_bytes() - def get_bytes( + def get_bytes_sync( self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None ) -> bytes: """ @@ -289,7 +289,7 @@ def get_bytes( Warnings -------- - Do not call this method from async functions. Use ``get_bytes_async()`` instead + Do not call this method from async functions. Use ``get_bytes()`` instead to avoid blocking the event loop. See Also @@ -300,23 +300,22 @@ def get_bytes( Examples -------- >>> store = MemoryStore() - >>> store.set("data", Buffer.from_bytes(b"hello world")) - >>> data = store.get_bytes("data", prototype=default_buffer_prototype()) + >>> await store.set("data", Buffer.from_bytes(b"hello world")) + >>> data = store.get_bytes_sync("data", prototype=default_buffer_prototype()) >>> print(data) b'hello world' """ - return sync(self.get_bytes_async(key, prototype=prototype, byte_range=byte_range)) + return sync(self.get_bytes(key, prototype=prototype, byte_range=byte_range)) - async def get_json_async( + async def get_json( self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None ) -> Any: """ Retrieve and parse JSON data from the store asynchronously. This is a convenience method that retrieves bytes from the store and - parses them as JSON. Commonly used for reading Zarr metadata files - like ``zarr.json``. + parses them as JSON. Parameters ---------- @@ -344,31 +343,29 @@ async def get_json_async( See Also -------- - get_bytes_async : Method for retrieving raw bytes without parsing. - get_json : Synchronous version of this method. + get_bytes : Method for retrieving raw bytes. + get_json_sync : Synchronous version of this method. Examples -------- >>> store = await MemoryStore.open() >>> metadata = {"zarr_format": 3, "node_type": "array"} >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) - >>> data = await store.get_json_async("zarr.json", prototype=default_buffer_prototype()) + >>> data = await store.get_json("zarr.json", prototype=default_buffer_prototype()) >>> print(data) {'zarr_format': 3, 'node_type': 'array'} """ - return json.loads( - await self.get_bytes_async(key, prototype=prototype, byte_range=byte_range) - ) + return json.loads(await self.get_bytes(key, prototype=prototype, byte_range=byte_range)) - def get_json( + def get_json_sync( self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None ) -> Any: """ Retrieve and parse JSON data from the store synchronously. - This is a synchronous wrapper around ``get_json_async()``. It should only - be called from non-async code. For async contexts, use ``get_json_async()`` + This is a synchronous wrapper around ``get_json()``. It should only + be called from non-async code. For async contexts, use ``get_json()`` instead. Parameters @@ -397,25 +394,25 @@ def get_json( Warnings -------- - Do not call this method from async functions. Use ``get_json_async()`` instead + Do not call this method from async functions. Use ``get_json()`` instead to avoid blocking the event loop. See Also -------- - get_json_async : Asynchronous version of this method. - get_bytes : Synchronous method for retrieving raw bytes without parsing. + get_json : Asynchronous version of this method. + get_bytes_sync : Synchronous method for retrieving raw bytes without parsing. Examples -------- >>> store = MemoryStore() >>> metadata = {"zarr_format": 3, "node_type": "array"} >>> store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) - >>> data = store.get_json("zarr.json", prototype=default_buffer_prototype()) + >>> data = store.get_json_sync("zarr.json", prototype=default_buffer_prototype()) >>> print(data) {'zarr_format': 3, 'node_type': 'array'} """ - return sync(self.get_json_async(key, prototype=prototype, byte_range=byte_range)) + return sync(self.get_json(key, prototype=prototype, byte_range=byte_range)) @abstractmethod async def get_partial_values( diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index 13c86a2f22..08681f2630 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -306,7 +306,7 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: except (FileNotFoundError, NotADirectoryError): pass - async def get_bytes_async( + async def get_bytes( self, key: str = "", *, @@ -356,9 +356,9 @@ async def get_bytes_async( """ if prototype is None: prototype = default_buffer_prototype() - return await super().get_bytes_async(key, prototype=prototype, byte_range=byte_range) + return await super().get_bytes(key, prototype=prototype, byte_range=byte_range) - def get_bytes( + def get_bytes_sync( self, key: str = "", *, @@ -412,9 +412,9 @@ def get_bytes( """ if prototype is None: prototype = default_buffer_prototype() - return super().get_bytes(key, prototype=prototype, byte_range=byte_range) + return super().get_bytes_sync(key, prototype=prototype, byte_range=byte_range) - async def get_json_async( + async def get_json( self, key: str = "", *, @@ -425,7 +425,7 @@ async def get_json_async( Retrieve and parse JSON data from the local store asynchronously. This is a convenience override that makes the ``prototype`` parameter optional - by defaulting to the standard buffer prototype. See the base ``Store.get_json_async`` + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` for full documentation. Parameters @@ -454,7 +454,7 @@ async def get_json_async( See Also -------- - Store.get_json_async : Base implementation with full documentation. + Store.get_json : Base implementation with full documentation. get_json : Synchronous version of this method. get_bytes_async : Method for retrieving raw bytes without parsing. @@ -465,15 +465,15 @@ async def get_json_async( >>> metadata = {"zarr_format": 3, "node_type": "array"} >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) >>> # No need to specify prototype for LocalStore - >>> data = await store.get_json_async("zarr.json") + >>> data = await store.get_json("zarr.json") >>> print(data) {'zarr_format': 3, 'node_type': 'array'} """ if prototype is None: prototype = default_buffer_prototype() - return await super().get_json_async(key, prototype=prototype, byte_range=byte_range) + return await super().get_json(key, prototype=prototype, byte_range=byte_range) - def get_json( + def get_json_sync( self, key: str = "", *, @@ -513,12 +513,12 @@ def get_json( Warnings -------- - Do not call this method from async functions. Use ``get_json_async()`` instead. + Do not call this method from async functions. Use ``get_json()`` instead. See Also -------- Store.get_json : Base implementation with full documentation. - get_json_async : Asynchronous version of this method. + get_json : Asynchronous version of this method. get_bytes : Method for retrieving raw bytes without parsing. Examples @@ -534,7 +534,7 @@ def get_json( """ if prototype is None: prototype = default_buffer_prototype() - return super().get_json(key, prototype=prototype, byte_range=byte_range) + return super().get_json_sync(key, prototype=prototype, byte_range=byte_range) async def move(self, dest_root: Path | str) -> None: """ diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index b56771f62a..5a2593eb25 100644 --- a/src/zarr/storage/_memory.py +++ b/src/zarr/storage/_memory.py @@ -175,7 +175,7 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]: for key in keys_unique: yield key - async def get_bytes_async( + async def get_bytes( self, key: str = "", *, @@ -225,9 +225,9 @@ async def get_bytes_async( """ if prototype is None: prototype = default_buffer_prototype() - return await super().get_bytes_async(key, prototype=prototype, byte_range=byte_range) + return await super().get_bytes(key, prototype=prototype, byte_range=byte_range) - def get_bytes( + def get_bytes_sync( self, key: str = "", *, @@ -281,9 +281,9 @@ def get_bytes( """ if prototype is None: prototype = default_buffer_prototype() - return super().get_bytes(key, prototype=prototype, byte_range=byte_range) + return super().get_bytes_sync(key, prototype=prototype, byte_range=byte_range) - async def get_json_async( + async def get_json( self, key: str = "", *, @@ -294,7 +294,7 @@ async def get_json_async( Retrieve and parse JSON data from the memory store asynchronously. This is a convenience override that makes the ``prototype`` parameter optional - by defaulting to the standard buffer prototype. See the base ``Store.get_json_async`` + by defaulting to the standard buffer prototype. See the base ``Store.get_json`` for full documentation. Parameters @@ -323,7 +323,7 @@ async def get_json_async( See Also -------- - Store.get_json_async : Base implementation with full documentation. + Store.get_json : Base implementation with full documentation. get_json : Synchronous version of this method. get_bytes_async : Method for retrieving raw bytes without parsing. @@ -334,15 +334,15 @@ async def get_json_async( >>> metadata = {"zarr_format": 3, "node_type": "array"} >>> await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode())) >>> # No need to specify prototype for MemoryStore - >>> data = await store.get_json_async("zarr.json") + >>> data = await store.get_json("zarr.json") >>> print(data) {'zarr_format': 3, 'node_type': 'array'} """ if prototype is None: prototype = default_buffer_prototype() - return await super().get_json_async(key, prototype=prototype, byte_range=byte_range) + return await super().get_json(key, prototype=prototype, byte_range=byte_range) - def get_json( + def get_json_sync( self, key: str = "", *, @@ -382,12 +382,12 @@ def get_json( Warnings -------- - Do not call this method from async functions. Use ``get_json_async()`` instead. + Do not call this method from async functions. Use ``get_json()`` instead. See Also -------- Store.get_json : Base implementation with full documentation. - get_json_async : Asynchronous version of this method. + get_json : Asynchronous version of this method. get_bytes : Method for retrieving raw bytes without parsing. Examples @@ -403,7 +403,7 @@ def get_json( """ if prototype is None: prototype = default_buffer_prototype() - return super().get_json(key, prototype=prototype, byte_range=byte_range) + return super().get_json_sync(key, prototype=prototype, byte_range=byte_range) class GpuMemoryStore(MemoryStore): diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index 30ff376fb0..f0cb6dd48f 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -534,9 +534,9 @@ async def test_get_bytes_async(self, store: S) -> None: data = b"hello world" key = "zarr.json" await self.set(store, key, self.buffer_cls.from_bytes(data)) - assert await store.get_bytes_async(key, prototype=default_buffer_prototype()) == data + assert await store.get_bytes(key, prototype=default_buffer_prototype()) == data with pytest.raises(FileNotFoundError): - await store.get_bytes_async("nonexistent_key", prototype=default_buffer_prototype()) + await store.get_bytes("nonexistent_key", prototype=default_buffer_prototype()) def test_get_bytes_sync(self, store: S) -> None: """ @@ -545,9 +545,9 @@ def test_get_bytes_sync(self, store: S) -> None: data = b"hello world" key = "zarr.json" sync(self.set(store, key, self.buffer_cls.from_bytes(data))) - assert store.get_bytes(key, prototype=default_buffer_prototype()) == data + assert store.get_bytes_sync(key, prototype=default_buffer_prototype()) == data - async def test_get_json_async(self, store: S) -> None: + async def test_get_json(self, store: S) -> None: """ Test that the get_bytes_async method reads json. """ @@ -555,7 +555,7 @@ async def test_get_json_async(self, store: S) -> None: data_bytes = json.dumps(data).encode("utf-8") key = "zarr.json" await self.set(store, key, self.buffer_cls.from_bytes(data_bytes)) - assert await store.get_json_async(key, prototype=default_buffer_prototype()) == data + assert await store.get_json(key, prototype=default_buffer_prototype()) == data def test_get_json_sync(self, store: S) -> None: """ @@ -565,7 +565,7 @@ def test_get_json_sync(self, store: S) -> None: data_bytes = json.dumps(data).encode("utf-8") key = "zarr.json" sync(self.set(store, key, self.buffer_cls.from_bytes(data_bytes))) - assert store.get_json(key, prototype=default_buffer_prototype()) == data + assert store.get_json_sync(key, prototype=default_buffer_prototype()) == data class LatencyStore(WrapperStore[Store]): diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index 35d48e3f95..d6a97110ad 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import pathlib import re @@ -9,6 +10,8 @@ import zarr from zarr import create_array from zarr.core.buffer import Buffer, cpu +from zarr.core.buffer.core import BufferPrototype, default_buffer_prototype +from zarr.core.sync import sync from zarr.storage import LocalStore from zarr.storage._local import _atomic_write from zarr.testing.store import StoreTests @@ -153,9 +156,8 @@ def test_atomic_write_exclusive_preexisting(tmp_path: pathlib.Path) -> None: async def test_get_bytes_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_bytes_async works with prototype=None.""" + """Test that get_bytes works with prototype=None.""" from zarr.core.buffer import cpu - from zarr.core.buffer.core import default_buffer_prototype store = await LocalStore.open(root=tmp_path) data = b"hello world" @@ -163,18 +165,17 @@ async def test_get_bytes_with_prototype_none(tmp_path: pathlib.Path) -> None: await store.set(key, cpu.Buffer.from_bytes(data)) # Test with None (default) - result_none = await store.get_bytes_async(key) + result_none = await store.get_bytes(key) assert result_none == data # Test with explicit prototype - result_proto = await store.get_bytes_async(key, prototype=default_buffer_prototype()) + result_proto = await store.get_bytes(key, prototype=default_buffer_prototype()) assert result_proto == data def test_get_bytes_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_bytes works with prototype=None.""" + """Test that get_bytes_sync works with prototype=None.""" from zarr.core.buffer import cpu - from zarr.core.buffer.core import default_buffer_prototype from zarr.core.sync import sync store = sync(LocalStore.open(root=tmp_path)) @@ -183,20 +184,19 @@ def test_get_bytes_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: sync(store.set(key, cpu.Buffer.from_bytes(data))) # Test with None (default) - result_none = store.get_bytes(key) + result_none = store.get_bytes_sync(key) assert result_none == data # Test with explicit prototype - result_proto = store.get_bytes(key, prototype=default_buffer_prototype()) + result_proto = store.get_bytes_sync(key, prototype=default_buffer_prototype()) assert result_proto == data -async def test_get_json_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_json_async works with prototype=None.""" - import json - - from zarr.core.buffer import cpu - from zarr.core.buffer.core import default_buffer_prototype +@pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) +async def test_get_json_with_prototype_none( + tmp_path: pathlib.Path, buffer_cls: None | BufferPrototype +) -> None: + """Test that get_json works with prototype=None.""" store = await LocalStore.open(root=tmp_path) data = {"foo": "bar", "number": 42} @@ -204,21 +204,15 @@ async def test_get_json_with_prototype_none(tmp_path: pathlib.Path) -> None: await store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode())) # Test with None (default) - result_none = await store.get_json_async(key) - assert result_none == data + result = await store.get_json(key, prototype=buffer_cls) + assert result == data - # Test with explicit prototype - result_proto = await store.get_json_async(key, prototype=default_buffer_prototype()) - assert result_proto == data - -def test_get_json_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_json works with prototype=None.""" - import json - - from zarr.core.buffer import cpu - from zarr.core.buffer.core import default_buffer_prototype - from zarr.core.sync import sync +@pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) +def test_get_json_sync_with_prototype( + tmp_path: pathlib.Path, buffer_cls: None | BufferPrototype +) -> None: + """Test that get_json_sync works with prototype=None.""" store = sync(LocalStore.open(root=tmp_path)) data = {"foo": "bar", "number": 42} @@ -226,9 +220,5 @@ def test_get_json_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: sync(store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode()))) # Test with None (default) - result_none = store.get_json(key) - assert result_none == data - - # Test with explicit prototype - result_proto = store.get_json(key, prototype=default_buffer_prototype()) - assert result_proto == data + result = store.get_json_sync(key, prototype=buffer_cls) + assert result == data diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index b56d9933d4..c47c1adb12 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -85,11 +85,11 @@ async def test_get_bytes_with_prototype_none(self, store: MemoryStore) -> None: await self.set(store, key, self.buffer_cls.from_bytes(data)) # Test with None (default) - result_none = await store.get_bytes_async(key) + result_none = await store.get_bytes(key) assert result_none == data # Test with explicit prototype - result_proto = await store.get_bytes_async(key, prototype=default_buffer_prototype()) + result_proto = await store.get_bytes(key, prototype=default_buffer_prototype()) assert result_proto == data def test_get_bytes_sync_with_prototype_none(self, store: MemoryStore) -> None: @@ -102,15 +102,15 @@ def test_get_bytes_sync_with_prototype_none(self, store: MemoryStore) -> None: sync(self.set(store, key, self.buffer_cls.from_bytes(data))) # Test with None (default) - result_none = store.get_bytes(key) + result_none = store.get_bytes_sync(key) assert result_none == data # Test with explicit prototype - result_proto = store.get_bytes(key, prototype=default_buffer_prototype()) + result_proto = store.get_bytes_sync(key, prototype=default_buffer_prototype()) assert result_proto == data async def test_get_json_with_prototype_none(self, store: MemoryStore) -> None: - """Test that get_json_async works with prototype=None.""" + """Test that get_json works with prototype=None.""" import json from zarr.core.buffer.core import default_buffer_prototype @@ -120,11 +120,11 @@ async def test_get_json_with_prototype_none(self, store: MemoryStore) -> None: await self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode())) # Test with None (default) - result_none = await store.get_json_async(key) + result_none = await store.get_json(key) assert result_none == data # Test with explicit prototype - result_proto = await store.get_json_async(key, prototype=default_buffer_prototype()) + result_proto = await store.get_json(key, prototype=default_buffer_prototype()) assert result_proto == data def test_get_json_sync_with_prototype_none(self, store: MemoryStore) -> None: @@ -139,11 +139,11 @@ def test_get_json_sync_with_prototype_none(self, store: MemoryStore) -> None: sync(self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode()))) # Test with None (default) - result_none = store.get_json(key) + result_none = store.get_json_sync(key) assert result_none == data # Test with explicit prototype - result_proto = store.get_json(key, prototype=default_buffer_prototype()) + result_proto = store.get_json_sync(key, prototype=default_buffer_prototype()) assert result_proto == data From 38ff5172cc126d15abfa992ceee9cba81f7e9e3e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 8 Jan 2026 11:02:07 +0100 Subject: [PATCH 6/7] continue renaming / test refactoring --- src/zarr/abc/store.py | 10 ++--- src/zarr/storage/_local.py | 22 ++++----- src/zarr/storage/_memory.py | 22 ++++----- src/zarr/testing/store.py | 8 ++-- tests/test_store/test_memory.py | 79 +++++++++++++-------------------- 5 files changed, 61 insertions(+), 80 deletions(-) diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index e685e4b3b0..0e98777ff5 100644 --- a/src/zarr/abc/store.py +++ b/src/zarr/abc/store.py @@ -248,7 +248,7 @@ async def get_bytes( -------- >>> store = await MemoryStore.open() >>> await store.set("data", Buffer.from_bytes(b"hello world")) - >>> data = await store.get_bytes_async("data", prototype=default_buffer_prototype()) + >>> data = await store.get_bytes("data", prototype=default_buffer_prototype()) >>> print(data) b'hello world' """ @@ -263,8 +263,8 @@ def get_bytes_sync( """ Retrieve raw bytes from the store synchronously. - This is a synchronous wrapper around ``get_bytes_async()``. It should only - be called from non-async code. For async contexts, use ``get_bytes_async()`` + This is a synchronous wrapper around ``get_bytes()``. It should only + be called from non-async code. For async contexts, use ``get_bytes()`` instead. Parameters @@ -294,8 +294,8 @@ def get_bytes_sync( See Also -------- - get_bytes_async : Asynchronous version of this method. - get_json : Synchronous method for retrieving and parsing JSON data. + get_bytes : Asynchronous version of this method. + get_json_sync : Synchronous method for retrieving and parsing JSON data. Examples -------- diff --git a/src/zarr/storage/_local.py b/src/zarr/storage/_local.py index 08681f2630..9fb3f8b6ad 100644 --- a/src/zarr/storage/_local.py +++ b/src/zarr/storage/_local.py @@ -317,7 +317,7 @@ async def get_bytes( Retrieve raw bytes from the local store asynchronously. This is a convenience override that makes the ``prototype`` parameter optional - by defaulting to the standard buffer prototype. See the base ``Store.get_bytes_async`` + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` for full documentation. Parameters @@ -342,15 +342,15 @@ async def get_bytes( See Also -------- - Store.get_bytes_async : Base implementation with full documentation. - get_bytes : Synchronous version of this method. + Store.get_bytes : Base implementation with full documentation. + get_bytes_sync : Synchronous version of this method. Examples -------- >>> store = await LocalStore.open("data") >>> await store.set("data", Buffer.from_bytes(b"hello")) >>> # No need to specify prototype for LocalStore - >>> data = await store.get_bytes_async("data") + >>> data = await store.get_bytes("data") >>> print(data) b'hello' """ @@ -394,12 +394,12 @@ def get_bytes_sync( Warnings -------- - Do not call this method from async functions. Use ``get_bytes_async()`` instead. + Do not call this method from async functions. Use ``get_bytes()`` instead. See Also -------- - Store.get_bytes : Base implementation with full documentation. - get_bytes_async : Asynchronous version of this method. + Store.get_bytes_sync : Base implementation with full documentation. + get_bytes : Asynchronous version of this method. Examples -------- @@ -455,8 +455,8 @@ async def get_json( See Also -------- Store.get_json : Base implementation with full documentation. - get_json : Synchronous version of this method. - get_bytes_async : Method for retrieving raw bytes without parsing. + get_json_sync : Synchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. Examples -------- @@ -517,9 +517,9 @@ def get_json_sync( See Also -------- - Store.get_json : Base implementation with full documentation. + Store.get_json_sync : Base implementation with full documentation. get_json : Asynchronous version of this method. - get_bytes : Method for retrieving raw bytes without parsing. + get_bytes_sync : Method for retrieving raw bytes without parsing. Examples -------- diff --git a/src/zarr/storage/_memory.py b/src/zarr/storage/_memory.py index 5a2593eb25..1568cc6736 100644 --- a/src/zarr/storage/_memory.py +++ b/src/zarr/storage/_memory.py @@ -186,7 +186,7 @@ async def get_bytes( Retrieve raw bytes from the memory store asynchronously. This is a convenience override that makes the ``prototype`` parameter optional - by defaulting to the standard buffer prototype. See the base ``Store.get_bytes_async`` + by defaulting to the standard buffer prototype. See the base ``Store.get_bytes`` for full documentation. Parameters @@ -211,15 +211,15 @@ async def get_bytes( See Also -------- - Store.get_bytes_async : Base implementation with full documentation. - get_bytes : Synchronous version of this method. + Store.get_bytes : Base implementation with full documentation. + get_bytes_sync : Synchronous version of this method. Examples -------- >>> store = await MemoryStore.open() >>> await store.set("data", Buffer.from_bytes(b"hello")) >>> # No need to specify prototype for MemoryStore - >>> data = await store.get_bytes_async("data") + >>> data = await store.get_bytes("data") >>> print(data) b'hello' """ @@ -263,12 +263,12 @@ def get_bytes_sync( Warnings -------- - Do not call this method from async functions. Use ``get_bytes_async()`` instead. + Do not call this method from async functions. Use ``get_bytes()`` instead. See Also -------- - Store.get_bytes : Base implementation with full documentation. - get_bytes_async : Asynchronous version of this method. + Store.get_bytes_sync : Base implementation with full documentation. + get_bytes : Asynchronous version of this method. Examples -------- @@ -324,8 +324,8 @@ async def get_json( See Also -------- Store.get_json : Base implementation with full documentation. - get_json : Synchronous version of this method. - get_bytes_async : Method for retrieving raw bytes without parsing. + get_json_sync : Synchronous version of this method. + get_bytes : Method for retrieving raw bytes without parsing. Examples -------- @@ -386,9 +386,9 @@ def get_json_sync( See Also -------- - Store.get_json : Base implementation with full documentation. + Store.get_json_sync : Base implementation with full documentation. get_json : Asynchronous version of this method. - get_bytes : Method for retrieving raw bytes without parsing. + get_bytes_sync : Method for retrieving raw bytes without parsing. Examples -------- diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index f0cb6dd48f..a56061ae12 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -527,9 +527,9 @@ async def test_set_if_not_exists(self, store: S) -> None: result = await store.get("k2", default_buffer_prototype()) assert result == new - async def test_get_bytes_async(self, store: S) -> None: + async def test_get_bytes(self, store: S) -> None: """ - Test that the get_bytes_async method reads bytes. + Test that the get_bytes method reads bytes. """ data = b"hello world" key = "zarr.json" @@ -540,7 +540,7 @@ async def test_get_bytes_async(self, store: S) -> None: def test_get_bytes_sync(self, store: S) -> None: """ - Test that the get_bytes method reads bytes. + Test that the get_bytes_sync method reads bytes. """ data = b"hello world" key = "zarr.json" @@ -549,7 +549,7 @@ def test_get_bytes_sync(self, store: S) -> None: async def test_get_json(self, store: S) -> None: """ - Test that the get_bytes_async method reads json. + Test that the get_json method reads json. """ data = {"foo": "bar"} data_bytes = json.dumps(data).encode("utf-8") diff --git a/tests/test_store/test_memory.py b/tests/test_store/test_memory.py index c47c1adb12..96b7fe9845 100644 --- a/tests/test_store/test_memory.py +++ b/tests/test_store/test_memory.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import re from typing import TYPE_CHECKING, Any @@ -9,12 +10,14 @@ import zarr from zarr.core.buffer import Buffer, cpu, gpu +from zarr.core.sync import sync from zarr.errors import ZarrUserWarning from zarr.storage import GpuMemoryStore, MemoryStore from zarr.testing.store import StoreTests from zarr.testing.utils import gpu_test if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype from zarr.core.common import ZarrFormat @@ -76,75 +79,53 @@ async def test_deterministic_size( np.testing.assert_array_equal(a[:3], 1) np.testing.assert_array_equal(a[3:], 0) - async def test_get_bytes_with_prototype_none(self, store: MemoryStore) -> None: - """Test that get_bytes_async works with prototype=None.""" - from zarr.core.buffer.core import default_buffer_prototype - + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + async def test_get_bytes_with_prototype_none( + self, store: MemoryStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_bytes works with prototype=None.""" data = b"hello world" key = "test_key" await self.set(store, key, self.buffer_cls.from_bytes(data)) - # Test with None (default) - result_none = await store.get_bytes(key) - assert result_none == data - - # Test with explicit prototype - result_proto = await store.get_bytes(key, prototype=default_buffer_prototype()) - assert result_proto == data - - def test_get_bytes_sync_with_prototype_none(self, store: MemoryStore) -> None: - """Test that get_bytes works with prototype=None.""" - from zarr.core.buffer.core import default_buffer_prototype - from zarr.core.sync import sync + result = await store.get_bytes(key, prototype=buffer_cls) + assert result == data + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + def test_get_bytes_sync_with_prototype_none( + self, store: MemoryStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_bytes_sync works with prototype=None.""" data = b"hello world" key = "test_key" sync(self.set(store, key, self.buffer_cls.from_bytes(data))) - # Test with None (default) - result_none = store.get_bytes_sync(key) - assert result_none == data + result = store.get_bytes_sync(key, prototype=buffer_cls) + assert result == data - # Test with explicit prototype - result_proto = store.get_bytes_sync(key, prototype=default_buffer_prototype()) - assert result_proto == data - - async def test_get_json_with_prototype_none(self, store: MemoryStore) -> None: + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + async def test_get_json_with_prototype_none( + self, store: MemoryStore, buffer_cls: None | BufferPrototype + ) -> None: """Test that get_json works with prototype=None.""" - import json - - from zarr.core.buffer.core import default_buffer_prototype - data = {"foo": "bar", "number": 42} key = "test.json" await self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode())) - # Test with None (default) - result_none = await store.get_json(key) - assert result_none == data - - # Test with explicit prototype - result_proto = await store.get_json(key, prototype=default_buffer_prototype()) - assert result_proto == data - - def test_get_json_sync_with_prototype_none(self, store: MemoryStore) -> None: - """Test that get_json works with prototype=None.""" - import json - - from zarr.core.buffer.core import default_buffer_prototype - from zarr.core.sync import sync + result = await store.get_json(key, prototype=buffer_cls) + assert result == data + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + def test_get_json_sync_with_prototype_none( + self, store: MemoryStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_json_sync works with prototype=None.""" data = {"foo": "bar", "number": 42} key = "test.json" sync(self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode()))) - # Test with None (default) - result_none = store.get_json_sync(key) - assert result_none == data - - # Test with explicit prototype - result_proto = store.get_json_sync(key, prototype=default_buffer_prototype()) - assert result_proto == data + result = store.get_json_sync(key, prototype=buffer_cls) + assert result == data # TODO: fix this warning From bdc4ef864b3bcbe422ab981eab6ec94f7af3ac0a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 8 Jan 2026 11:47:14 +0100 Subject: [PATCH 7/7] refactor new test functions --- tests/test_store/test_local.py | 122 ++++++++++++++------------------- 1 file changed, 52 insertions(+), 70 deletions(-) diff --git a/tests/test_store/test_local.py b/tests/test_store/test_local.py index d6a97110ad..fa4bc7cfc0 100644 --- a/tests/test_store/test_local.py +++ b/tests/test_store/test_local.py @@ -3,6 +3,7 @@ import json import pathlib import re +from typing import TYPE_CHECKING import numpy as np import pytest @@ -10,13 +11,15 @@ import zarr from zarr import create_array from zarr.core.buffer import Buffer, cpu -from zarr.core.buffer.core import BufferPrototype, default_buffer_prototype from zarr.core.sync import sync from zarr.storage import LocalStore from zarr.storage._local import _atomic_write from zarr.testing.store import StoreTests from zarr.testing.utils import assert_bytes_equal +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype + class TestLocalStore(StoreTests[LocalStore, cpu.Buffer]): store_cls = LocalStore @@ -111,6 +114,54 @@ async def test_move( ): await store2.move(destination) + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + async def test_get_bytes_with_prototype_none( + self, store: LocalStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_bytes works with prototype=None.""" + data = b"hello world" + key = "test_key" + await self.set(store, key, self.buffer_cls.from_bytes(data)) + + result = await store.get_bytes(key, prototype=buffer_cls) + assert result == data + + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + def test_get_bytes_sync_with_prototype_none( + self, store: LocalStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_bytes_sync works with prototype=None.""" + data = b"hello world" + key = "test_key" + sync(self.set(store, key, self.buffer_cls.from_bytes(data))) + + result = store.get_bytes_sync(key, prototype=buffer_cls) + assert result == data + + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + async def test_get_json_with_prototype_none( + self, store: LocalStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_json works with prototype=None.""" + data = {"foo": "bar", "number": 42} + key = "test.json" + await self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode())) + + result = await store.get_json(key, prototype=buffer_cls) + assert result == data + + @pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) + def test_get_json_sync_with_prototype_none( + self, store: LocalStore, buffer_cls: None | BufferPrototype + ) -> None: + """Test that get_json_sync works with prototype=None.""" + data = {"foo": "bar", "number": 42} + key = "test.json" + sync(self.set(store, key, self.buffer_cls.from_bytes(json.dumps(data).encode()))) + + result = store.get_json_sync(key, prototype=buffer_cls) + assert result == data + @pytest.mark.parametrize("exclusive", [True, False]) def test_atomic_write_successful(tmp_path: pathlib.Path, exclusive: bool) -> None: @@ -153,72 +204,3 @@ def test_atomic_write_exclusive_preexisting(tmp_path: pathlib.Path) -> None: f.write(b"abc") assert path.read_bytes() == b"xyz" assert list(path.parent.iterdir()) == [path] # no temp files - - -async def test_get_bytes_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_bytes works with prototype=None.""" - from zarr.core.buffer import cpu - - store = await LocalStore.open(root=tmp_path) - data = b"hello world" - key = "test_key" - await store.set(key, cpu.Buffer.from_bytes(data)) - - # Test with None (default) - result_none = await store.get_bytes(key) - assert result_none == data - - # Test with explicit prototype - result_proto = await store.get_bytes(key, prototype=default_buffer_prototype()) - assert result_proto == data - - -def test_get_bytes_sync_with_prototype_none(tmp_path: pathlib.Path) -> None: - """Test that get_bytes_sync works with prototype=None.""" - from zarr.core.buffer import cpu - from zarr.core.sync import sync - - store = sync(LocalStore.open(root=tmp_path)) - data = b"hello world" - key = "test_key" - sync(store.set(key, cpu.Buffer.from_bytes(data))) - - # Test with None (default) - result_none = store.get_bytes_sync(key) - assert result_none == data - - # Test with explicit prototype - result_proto = store.get_bytes_sync(key, prototype=default_buffer_prototype()) - assert result_proto == data - - -@pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) -async def test_get_json_with_prototype_none( - tmp_path: pathlib.Path, buffer_cls: None | BufferPrototype -) -> None: - """Test that get_json works with prototype=None.""" - - store = await LocalStore.open(root=tmp_path) - data = {"foo": "bar", "number": 42} - key = "test.json" - await store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode())) - - # Test with None (default) - result = await store.get_json(key, prototype=buffer_cls) - assert result == data - - -@pytest.mark.parametrize("buffer_cls", [None, cpu.buffer_prototype]) -def test_get_json_sync_with_prototype( - tmp_path: pathlib.Path, buffer_cls: None | BufferPrototype -) -> None: - """Test that get_json_sync works with prototype=None.""" - - store = sync(LocalStore.open(root=tmp_path)) - data = {"foo": "bar", "number": 42} - key = "test.json" - sync(store.set(key, cpu.Buffer.from_bytes(json.dumps(data).encode()))) - - # Test with None (default) - result = store.get_json_sync(key, prototype=buffer_cls) - assert result == data