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