turbo_broccoli.custom.bytes

 1"""bytes (de)serialization utilities."""
 2
 3from base64 import b64decode, b64encode
 4from math import ceil
 5from typing import Any
 6
 7from ..context import Context
 8from ..exceptions import DeserializationError, TypeNotSupported
 9
10
11def _bytes_from_json_v3(dct: dict, ctx: Context) -> bytes:
12    if "data" in dct:
13        return b64decode(dct["data"])
14    path = ctx.id_to_artifact_path(dct["id"])
15    with path.open(mode="rb") as fp:
16        return fp.read()
17
18
19def from_json(dct: dict, ctx: Context) -> bytes | None:
20    decoders = {
21        3: _bytes_from_json_v3,
22    }
23    try:
24        return decoders[dct["__version__"]](dct, ctx)
25    except KeyError as exc:
26        raise DeserializationError() from exc
27
28
29def to_json(obj: Any, ctx: Context) -> dict:
30    """
31    Serializes a Python `bytes` object into JSON using a base64 + ASCII
32    scheme. The return dict has the following structure
33
34    ```py
35    {
36        "__type__": "bytes",
37        "__version__": 3,
38        "data": <ASCII str>,
39    }
40    ```
41
42    or
43
44    ```py
45    {
46        "__type__": "bytes",
47        "__version__": 3,
48        "id": <uuid4>,
49    }
50    ```
51
52    if the base64 encoding of the object is too large.
53
54    """
55    if not isinstance(obj, bytes):
56        raise TypeNotSupported()
57    # https://stackoverflow.com/a/32140193
58    b64_size = (ceil((len(obj) * 4) / 3) + 3) & ~3
59    if b64_size <= ctx.min_artifact_size:
60        return {
61            "__type__": "bytes",
62            "__version__": 3,
63            "data": b64encode(obj).decode("ascii"),
64        }
65    path, name = ctx.new_artifact_path()
66    with path.open(mode="wb") as fp:
67        fp.write(obj)
68    return {"__type__": "bytes", "__version__": 3, "id": name}