Skip to content

Commit 4bfe00b

Browse files
d-v-bclaudechuckwondo
authored
refactor/consolidated JSON IO routines (#3998)
* feat: add free functions for JSON document I/O Add four free functions for moving JSON documents in and out of stores, plus a thin StorePath.get_json wrapper: - buffer_to_json / json_to_buffer: convert between a Buffer and a parsed JSON value. The buffer prototype and the serialization policy (indent from config, allow_nan=True) live here, in one place. - get_json / set_json: compose those with Store.get / Store.set. get_json returns None for a missing key, matching what most callers want; callers needing presence check for None themselves. These are free functions, not Store ABC methods, so stores cannot (and need not) override them and the Store contract gains no dependency on the buffer prototype or global config. Subsequent commits sweep the hand-rolled JSON I/O sites onto these and delete the old private _get_bytes/_get_json store methods. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * refactor: route JSON document I/O through the free functions Sweep the hand-rolled `json.dumps`/`json.loads` + buffer construction in the metadata write paths and the metadata/node read paths onto the free functions added in the previous commit. Write sites (`to_buffer_dict` in group/metadata) now call `json_to_buffer`, which centralizes the `json_indent`/`allow_nan=True` serialization policy. Read sites parse buffers through `buffer_to_json_object`, a new helper that narrows the parsed `JSON` to `dict[str, JSON]` once (every metadata document is an object) rather than relying on `json.loads` returning `Any`. The compact (no-indent) consolidated-metadata blob in `GroupMetadata.to_buffer_dict` is intentionally left as a raw `json.dumps` to preserve its byte layout, and `_read_metadata_v2` keeps its `asyncio.gather` of three `store.get` calls. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor: remove the unused private JSON/bytes store methods `Store._get_bytes`, `_get_bytes_sync`, `_get_json`, and `_get_json_sync` had no production callers — the metadata read/write paths now go through the free functions in `zarr.core._json`. Delete the four ABC methods, the eight per-store overrides on `LocalStore`/`MemoryStore` (which existed only to make the `prototype` argument optional), and the tests that exercised them (the shared `StoreTests` methods plus the per-store prototype=None tests). Drop the now-unused `json`/`sync`/`Any` imports they pulled in. These methods were always private; removing them is not a public API change. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * doc: add changelog fragment for the JSON I/O refactor Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor: make json_to_buffer pure (no config read) `json_to_buffer` (and `set_json`) now take the JSON encoding parameters `indent` and `allow_nan` as explicit keyword arguments instead of reading `config.get("json_indent")` internally. The functions in `zarr.core._json` are now pure: they depend only on their arguments, not the global config. Each `to_buffer_dict` caller (group, metadata v2/v3) reads `config.get("json_indent")` and passes it as `indent=`. This re-localizes the config read to the metadata layer that owns the policy, and keeps the I/O helpers free of any config dependency. With `indent` now explicit, the compact consolidated-metadata blob in `GroupMetadata.to_buffer_dict` also routes through `json_to_buffer` (indent defaults to None == compact), removing the last raw `json.dumps` in that method and the `json` import from group.py. Output bytes are unchanged (the consolidated/metadata regression tests pass). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * doc: rename changelog fragment to PR number 3998 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * test: cover the malformed-metadata branches in contains_array/group The except clauses in `contains_array`, `contains_group`, and `_contains_node_v3` (widened to catch `TypeError` in the JSON I/O refactor) had no test exercising them — a stored `zarr.json` that is malformed bytes, valid-but-non-object JSON, or an object missing `node_type` now has a test asserting each function reports the node as absent. One test per failure mode; the two `contains_*` functions are parametrized over `func`. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * Update src/zarr/core/_json.py Co-authored-by: Chuck Daniels <cjdaniels4@gmail.com> * fix: missing import * doc: address review nits on _json.py docstrings Add a Parameters section to buffer_to_json_object and drop the redundant type annotations from the numpydoc Parameters sections, per review feedback on #3998. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Chuck Daniels <cjdaniels4@gmail.com>
1 parent 467dda7 commit 4bfe00b

15 files changed

Lines changed: 378 additions & 967 deletions

File tree

changes/3998.misc.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Centralized JSON document I/O behind free functions in `zarr.core._json` and removed the unused private `Store._get_bytes`/`_get_json` methods and their per-store overrides.

src/zarr/abc/store.py

Lines changed: 0 additions & 225 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,12 @@
11
from __future__ import annotations
22

33
import asyncio
4-
import json
54
from abc import ABC, abstractmethod
65
from dataclasses import dataclass
76
from functools import partial
87
from itertools import starmap
98
from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable
109

11-
from zarr.core.sync import sync
12-
1310
if TYPE_CHECKING:
1411
from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Sequence
1512
from types import TracebackType
@@ -219,228 +216,6 @@ async def get(
219216
"""
220217
...
221218

222-
async def _get_bytes(
223-
self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
224-
) -> bytes:
225-
"""
226-
Retrieve raw bytes from the store asynchronously.
227-
228-
This is a convenience method that wraps ``get()`` and converts the result
229-
to bytes. Use this when you need the raw byte content of a stored value.
230-
231-
Parameters
232-
----------
233-
key : str
234-
The key identifying the data to retrieve.
235-
prototype : BufferPrototype
236-
The buffer prototype to use for reading the data.
237-
byte_range : ByteRequest, optional
238-
If specified, only retrieve a portion of the stored data.
239-
Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
240-
241-
Returns
242-
-------
243-
bytes
244-
The raw bytes stored at the given key.
245-
246-
Raises
247-
------
248-
FileNotFoundError
249-
If the key does not exist in the store.
250-
251-
See Also
252-
--------
253-
get : Lower-level method that returns a Buffer object.
254-
get_bytes : Synchronous version of this method.
255-
get_json : Asynchronous method for retrieving and parsing JSON data.
256-
257-
Examples
258-
--------
259-
>>> async def example():
260-
... from zarr.core.buffer.cpu import Buffer
261-
... from zarr.storage import MemoryStore
262-
...
263-
... store = await MemoryStore.open()
264-
... await store.set("data", Buffer.from_bytes(b"hello world"))
265-
... # No need to specify prototype for MemoryStore
266-
... return await store._get_bytes("data")
267-
268-
>>> import asyncio
269-
>>> asyncio.run(example())
270-
b'hello world'
271-
"""
272-
273-
buffer = await self.get(key, prototype, byte_range)
274-
if buffer is None:
275-
raise FileNotFoundError(key)
276-
return buffer.to_bytes()
277-
278-
def _get_bytes_sync(
279-
self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
280-
) -> bytes:
281-
"""
282-
Retrieve raw bytes from the store synchronously.
283-
284-
This is a synchronous wrapper around ``get_bytes()``. It should only
285-
be called from non-async code. For async contexts, use ``get_bytes()``
286-
instead.
287-
288-
Parameters
289-
----------
290-
key : str, optional
291-
The key identifying the data to retrieve. Defaults to an empty string.
292-
prototype : BufferPrototype
293-
The buffer prototype to use for reading the data.
294-
byte_range : ByteRequest, optional
295-
If specified, only retrieve a portion of the stored data.
296-
Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
297-
298-
Returns
299-
-------
300-
bytes
301-
The raw bytes stored at the given key.
302-
303-
Raises
304-
------
305-
FileNotFoundError
306-
If the key does not exist in the store.
307-
308-
Warnings
309-
--------
310-
Do not call this method from async functions. Use ``get_bytes()`` instead
311-
to avoid blocking the event loop.
312-
313-
See Also
314-
--------
315-
get_bytes : Asynchronous version of this method.
316-
get_json_sync : Synchronous method for retrieving and parsing JSON data.
317-
318-
Examples
319-
--------
320-
>>> from zarr.core.buffer.cpu import Buffer
321-
>>> from zarr.storage import MemoryStore
322-
>>> store = MemoryStore()
323-
>>> store.set_sync("data", Buffer.from_bytes(b"hello world"))
324-
>>> store._get_bytes_sync("data") # No need to specify prototype for MemoryStore
325-
b'hello world'
326-
"""
327-
328-
return sync(self._get_bytes(key, prototype=prototype, byte_range=byte_range))
329-
330-
async def _get_json(
331-
self, key: str, *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
332-
) -> Any:
333-
"""
334-
Retrieve and parse JSON data from the store asynchronously.
335-
336-
This is a convenience method that retrieves bytes from the store and
337-
parses them as JSON.
338-
339-
Parameters
340-
----------
341-
key : str
342-
The key identifying the JSON data to retrieve.
343-
prototype : BufferPrototype
344-
The buffer prototype to use for reading the data.
345-
byte_range : ByteRequest, optional
346-
If specified, only retrieve a portion of the stored data.
347-
Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
348-
Note: Using byte ranges with JSON may result in invalid JSON.
349-
350-
Returns
351-
-------
352-
Any
353-
The parsed JSON data. This follows the behavior of ``json.loads()`` and
354-
can be any JSON-serializable type: dict, list, str, int, float, bool, or None.
355-
356-
Raises
357-
------
358-
FileNotFoundError
359-
If the key does not exist in the store.
360-
json.JSONDecodeError
361-
If the stored data is not valid JSON.
362-
363-
See Also
364-
--------
365-
get_bytes : Method for retrieving raw bytes.
366-
get_json_sync : Synchronous version of this method.
367-
368-
Examples
369-
--------
370-
>>> async def example():
371-
... from zarr.core.buffer.cpu import Buffer
372-
... from zarr.storage import MemoryStore
373-
...
374-
... store = await MemoryStore.open()
375-
... metadata = {"zarr_format": 3, "node_type": "array"}
376-
... await store.set("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode()))
377-
... # No need to specify prototype for MemoryStore
378-
... return await store._get_json("zarr.json")
379-
380-
>>> import asyncio
381-
>>> asyncio.run(example())
382-
{'zarr_format': 3, 'node_type': 'array'}
383-
"""
384-
385-
return json.loads(await self._get_bytes(key, prototype=prototype, byte_range=byte_range))
386-
387-
def _get_json_sync(
388-
self, key: str = "", *, prototype: BufferPrototype, byte_range: ByteRequest | None = None
389-
) -> Any:
390-
"""
391-
Retrieve and parse JSON data from the store synchronously.
392-
393-
This is a synchronous wrapper around ``get_json()``. It should only
394-
be called from non-async code. For async contexts, use ``get_json()``
395-
instead.
396-
397-
Parameters
398-
----------
399-
key : str, optional
400-
The key identifying the JSON data to retrieve. Defaults to an empty string.
401-
prototype : BufferPrototype
402-
The buffer prototype to use for reading the data.
403-
byte_range : ByteRequest, optional
404-
If specified, only retrieve a portion of the stored data.
405-
Can be a ``RangeByteRequest``, ``OffsetByteRequest``, or ``SuffixByteRequest``.
406-
Note: Using byte ranges with JSON may result in invalid JSON.
407-
408-
Returns
409-
-------
410-
Any
411-
The parsed JSON data. This follows the behavior of ``json.loads()`` and
412-
can be any JSON-serializable type: dict, list, str, int, float, bool, or None.
413-
414-
Raises
415-
------
416-
FileNotFoundError
417-
If the key does not exist in the store.
418-
json.JSONDecodeError
419-
If the stored data is not valid JSON.
420-
421-
Warnings
422-
--------
423-
Do not call this method from async functions. Use ``get_json()`` instead
424-
to avoid blocking the event loop.
425-
426-
See Also
427-
--------
428-
get_json : Asynchronous version of this method.
429-
get_bytes_sync : Synchronous method for retrieving raw bytes without parsing.
430-
431-
Examples
432-
--------
433-
>>> from zarr.core.buffer.cpu import Buffer
434-
>>> from zarr.storage import MemoryStore
435-
>>> store = MemoryStore()
436-
>>> metadata = {"zarr_format": 3, "node_type": "array"}
437-
>>> store.set_sync("zarr.json", Buffer.from_bytes(json.dumps(metadata).encode()))
438-
>>> store._get_json_sync("zarr.json") # No need to specify prototype for MemoryStore
439-
{'zarr_format': 3, 'node_type': 'array'}
440-
"""
441-
442-
return sync(self._get_json(key, prototype=prototype, byte_range=byte_range))
443-
444219
@abstractmethod
445220
async def get_partial_values(
446221
self,

src/zarr/core/_json.py

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
"""Helpers for moving JSON documents in and out of zarr stores.
2+
3+
These are free functions, deliberately not methods on the ``Store`` ABC:
4+
reading and writing JSON is a composition of the store's ``get``/``set``
5+
primitives with a buffer/JSON conversion, not part of the store contract.
6+
Keeping them as functions means stores cannot (and need not) override them,
7+
and the ``Store`` definition stays free of any dependency on the buffer
8+
prototype.
9+
10+
These functions are pure: the JSON encoding parameters (``indent``,
11+
``allow_nan``) are explicit arguments rather than read from the global config.
12+
Callers that want zarr's configured indentation pass
13+
``indent=config.get("json_indent")``.
14+
15+
Two layers:
16+
17+
- ``buffer_to_json`` / ``json_to_buffer`` convert between a ``Buffer`` and a
18+
parsed JSON value. The buffer prototype lives here, at buffer construction,
19+
where it is meaningful.
20+
- ``get_json`` / ``set_json`` compose those with ``Store.get`` / ``Store.set``.
21+
``get_json`` returns ``None`` for a missing key (the contract most callers
22+
want); callers that require presence check for ``None`` themselves.
23+
"""
24+
25+
from __future__ import annotations
26+
27+
import json
28+
from typing import TYPE_CHECKING, cast
29+
30+
from zarr.core.buffer import default_buffer_prototype
31+
32+
if TYPE_CHECKING:
33+
from zarr.abc.store import ByteRequest, Store
34+
from zarr.core.buffer import Buffer, BufferPrototype
35+
from zarr.core.common import JSON
36+
37+
38+
def buffer_to_json(buffer: Buffer) -> JSON:
39+
"""Parse the contents of a `Buffer` as a JSON value."""
40+
# json.loads is typed as returning Any; the result is by definition JSON.
41+
return cast("JSON", json.loads(buffer.to_bytes()))
42+
43+
44+
def buffer_to_json_object(buffer: Buffer) -> dict[str, JSON]:
45+
"""Parse the contents of a `Buffer` as a JSON object (a `dict`).
46+
47+
Every metadata document zarr reads is a JSON object, so this narrows the
48+
`JSON` union to `dict[str, JSON]` once, here, instead of at each call site.
49+
50+
Parameters
51+
----------
52+
buffer
53+
The buffer whose contents are parsed as a JSON object.
54+
55+
Raises
56+
------
57+
TypeError
58+
If the parsed value is not a JSON object.
59+
"""
60+
obj = buffer_to_json(buffer)
61+
if not isinstance(obj, dict):
62+
raise TypeError(f"Expected a JSON object, got {type(obj).__name__}.")
63+
return obj
64+
65+
66+
def json_to_buffer(
67+
obj: JSON,
68+
*,
69+
prototype: BufferPrototype | None = None,
70+
indent: int | None = None,
71+
allow_nan: bool = True,
72+
) -> Buffer:
73+
"""Serialize a JSON value into a `Buffer`.
74+
75+
Parameters
76+
----------
77+
obj
78+
The JSON-serializable value to encode.
79+
prototype
80+
The buffer prototype to construct the result with. Defaults to
81+
`default_buffer_prototype()`.
82+
indent
83+
Indentation passed to `json.dumps`. `None` (the default) writes
84+
without newline indentation, using json's default separators.
85+
Callers that want zarr's configured indentation pass
86+
`indent=config.get("json_indent")`.
87+
allow_nan
88+
Whether to permit `NaN`/`Infinity` in the output, passed to
89+
`json.dumps`.
90+
"""
91+
if prototype is None:
92+
prototype = default_buffer_prototype()
93+
return prototype.buffer.from_bytes(json.dumps(obj, indent=indent, allow_nan=allow_nan).encode())
94+
95+
96+
async def get_json(store: Store, key: str, *, byte_range: ByteRequest | None = None) -> JSON | None:
97+
"""Read and parse the JSON document at `key`, or `None` if it is absent.
98+
99+
Parameters
100+
----------
101+
store
102+
The store to read from.
103+
key
104+
The key identifying the JSON document.
105+
byte_range
106+
If given, read only this portion of the value. Note that a partial
107+
read of a JSON document may not be valid JSON.
108+
109+
Returns
110+
-------
111+
JSON or None
112+
The parsed JSON value, or `None` if `key` does not exist.
113+
"""
114+
buffer = await store.get(key, default_buffer_prototype(), byte_range)
115+
return None if buffer is None else buffer_to_json(buffer)
116+
117+
118+
async def set_json(
119+
store: Store,
120+
key: str,
121+
obj: JSON,
122+
*,
123+
prototype: BufferPrototype | None = None,
124+
indent: int | None = None,
125+
allow_nan: bool = True,
126+
) -> None:
127+
"""Serialize `obj` as JSON and write it to `key` in `store`.
128+
129+
`indent` and `allow_nan` are forwarded to `json_to_buffer`.
130+
"""
131+
await store.set(
132+
key, json_to_buffer(obj, prototype=prototype, indent=indent, allow_nan=allow_nan)
133+
)

0 commit comments

Comments
 (0)