turbo_broccoli.custom.embedded

  1"""
  2Embedded JSON documents.
  3
  4Serializing a `EmbeddedDict` or a `EmbeddedList` will (unconditionally) result
  5in its own JSON artefact being created and referenced by the main JSON
  6document.
  7"""
  8
  9from pathlib import Path
 10from typing import Any, Callable, Tuple
 11
 12from ..context import Context
 13from ..exceptions import DeserializationError, TypeNotSupported
 14
 15
 16class EmbeddedDict(dict):
 17    """See module documentation"""
 18
 19    _tb_artifact_id: str | None = None
 20
 21
 22class EmbeddedList(list):
 23    """See module documentation"""
 24
 25    _tb_artifact_id: str | None = None
 26
 27
 28def _get_artifact_path(
 29    obj: EmbeddedDict | EmbeddedList, ctx: Context
 30) -> tuple[Path, str]:
 31    """
 32    If the object has an `_tb_artifact_id`, return the corresponding path and
 33    id. Otherwise, generate new ones.
 34
 35    Warning:
 36        Does not NOT set the `_tb_artifact_id` attribute.
 37    """
 38    if obj._tb_artifact_id is None:
 39        return ctx.new_artifact_path(extension="json")
 40    name = obj._tb_artifact_id
 41    path = ctx.id_to_artifact_path(name, extension="json")
 42    return path, name
 43
 44
 45def _embedded_dict_to_json(obj: EmbeddedDict, ctx: Context) -> dict:
 46    from turbo_broccoli.turbo_broccoli import to_json as _to_json
 47
 48    path, name = _get_artifact_path(obj, ctx)
 49    with path.open("w", encoding="utf-8") as fp:
 50        fp.write(_to_json(dict(obj), ctx))
 51    obj._tb_artifact_id = name
 52    return {"__type__": "embedded.dict", "__version__": 1, "id": name}
 53
 54
 55# TODO: deduplicate with _embedded_dict_to_json
 56def _embedded_list_to_json(obj: EmbeddedList, ctx: Context) -> dict:
 57    from turbo_broccoli.turbo_broccoli import to_json as _to_json
 58
 59    path, name = _get_artifact_path(obj, ctx)
 60    with path.open("w", encoding="utf-8") as fp:
 61        fp.write(_to_json(list(obj), ctx))
 62    obj._tb_artifact_id = name
 63    return {"__type__": "embedded.list", "__version__": 1, "id": name}
 64
 65
 66def _json_to_embedded_dict(dct: dict, ctx: Context) -> EmbeddedDict:
 67    decoders = {
 68        1: _json_to_embedded_dict_v1,
 69    }
 70    return decoders[dct["__version__"]](dct, ctx)
 71
 72
 73def _json_to_embedded_dict_v1(dct: dict, ctx: Context) -> EmbeddedDict:
 74    from turbo_broccoli.turbo_broccoli import from_json as _from_json
 75
 76    path = ctx.id_to_artifact_path(dct["id"], extension="json")
 77    with path.open("r", encoding="utf-8") as fp:
 78        obj = EmbeddedDict(_from_json(fp.read(), ctx))
 79    obj._tb_artifact_id = dct["id"]
 80    return obj
 81
 82
 83def _json_to_embedded_list(dct: dict, ctx: Context) -> EmbeddedList:
 84    decoders = {
 85        1: _json_to_embedded_list_v1,
 86    }
 87    return decoders[dct["__version__"]](dct, ctx)
 88
 89
 90# TODO: deduplicate with _json_to_embedded_dict_v1
 91def _json_to_embedded_list_v1(dct: dict, ctx: Context) -> EmbeddedList:
 92    from turbo_broccoli.turbo_broccoli import from_json as _from_json
 93
 94    path = ctx.id_to_artifact_path(dct["id"], extension="json")
 95    with path.open("r", encoding="utf-8") as fp:
 96        obj = EmbeddedList(_from_json(fp.read(), ctx))
 97    obj._tb_artifact_id = dct["id"]
 98    return obj
 99
100
101def from_json(dct: dict, ctx: Context) -> Any:
102    decoders = {
103        "embedded.dict": _json_to_embedded_dict,
104        "embedded.list": _json_to_embedded_list,
105    }
106    try:
107        type_name = dct["__type__"]
108        return decoders[type_name](dct, ctx)
109    except KeyError as exc:
110        raise DeserializationError() from exc
111
112
113def to_json(obj: Any, ctx: Context) -> dict:
114    """
115    Serializes a `EmbeddedDict` or an `EmbeddedList` into JSON. The return dict
116    has the following structure
117
118    ```py
119    {
120        "__type__": "embedded.dict",
121        "__version__": 1,
122        "id": <uuid4>,
123    }
124    ```
125
126    or
127
128    ```py
129    {
130        "__type__": "embedded.list",
131        "__version__": 1,
132        "id": <uuid4>,
133    }
134    ```
135
136    where the UUID points to the artefact containing the actual data.
137    """
138    encoders: list[Tuple[type, Callable[[Any, Context], dict]]] = [
139        (EmbeddedDict, _embedded_dict_to_json),
140        (EmbeddedList, _embedded_list_to_json),
141    ]
142    for t, f in encoders:
143        if isinstance(obj, t):
144            return f(obj, ctx)
145    raise TypeNotSupported()