turbo_broccoli.custom.collections

  1"""Python standard collections and container types (de)serialization"""
  2
  3from collections import deque, namedtuple
  4from typing import Any, Callable, Tuple
  5
  6from ..context import Context
  7from ..exceptions import DeserializationError, TypeNotSupported
  8
  9
 10def _deque_to_json(deq: deque, ctx: Context) -> dict:
 11    return {
 12        "__type__": "collections.deque",
 13        "__version__": 2,
 14        "data": list(deq),
 15        "maxlen": deq.maxlen,
 16    }
 17
 18
 19def _json_to_deque(dct: dict, ctx: Context) -> deque | None:
 20    decoders = {
 21        2: _json_to_deque_v2,
 22    }
 23    return decoders[dct["__version__"]](dct, ctx)
 24
 25
 26def _json_to_deque_v2(dct: dict, ctx: Context) -> Any:
 27    return deque(dct["data"], dct["maxlen"])
 28
 29
 30def _json_to_namedtuple(dct: dict, ctx: Context) -> Any:
 31    decoders = {
 32        2: _json_to_namedtuple_v2,
 33    }
 34    return decoders[dct["__version__"]](dct, ctx)
 35
 36
 37def _json_to_namedtuple_v2(dct: dict, ctx: Context) -> Any:
 38    return namedtuple(dct["class"], dct["data"].keys())(**dct["data"])
 39
 40
 41def _json_to_set(dct: dict, ctx: Context) -> set:
 42    decoders = {
 43        2: _json_to_set_v2,
 44    }
 45    return decoders[dct["__version__"]](dct, ctx)
 46
 47
 48def _json_to_set_v2(dct: dict, ctx: Context) -> Any:
 49    return set(dct["data"])
 50
 51
 52def _json_to_tuple(dct: dict, ctx: Context) -> tuple:
 53    decoders = {
 54        1: _json_to_tuple_v1,
 55    }
 56    return decoders[dct["__version__"]](dct, ctx)
 57
 58
 59def _json_to_tuple_v1(dct: dict, ctx: Context) -> Any:
 60    return tuple(dct["data"])
 61
 62
 63def _set_to_json(obj: set, ctx: Context) -> dict:
 64    return {"__type__": "collections.set", "__version__": 2, "data": list(obj)}
 65
 66
 67def _tuple_to_json(obj: tuple, ctx: Context) -> dict:
 68    """
 69    Converts a tuple or namedtuple into a JSON document.
 70
 71    A tuple is a namedtuple if it has the following attributes: `_asdict`,
 72    `_field_defaults`, `_fields`, `_make`, `_replace`. See
 73    https://docs.python.org/3/library/collections.html#collections.namedtuple .
 74    """
 75    attributes = ["_asdict", "_field_defaults", "_fields", "_make", "_replace"]
 76    if not all(map(lambda a: hasattr(obj, a), attributes)):
 77        return {
 78            "__type__": "collections.tuple",
 79            "__version__": 1,
 80            "data": list(obj),
 81        }
 82    return {
 83        "__type__": "collections.namedtuple",
 84        "__version__": 2,
 85        "class": obj.__class__.__name__,
 86        "data": obj._asdict(),  # type: ignore
 87    }
 88
 89
 90def from_json(dct: dict, ctx: Context) -> Any:
 91    decoders = {
 92        "collections.deque": _json_to_deque,
 93        "collections.namedtuple": _json_to_namedtuple,
 94        "collections.set": _json_to_set,
 95        "collections.tuple": _json_to_tuple,
 96    }
 97    try:
 98        type_name = dct["__type__"]
 99        return decoders[type_name](dct, ctx)
100    except KeyError as exc:
101        raise DeserializationError() from exc
102
103
104def to_json(obj: Any, ctx: Context) -> dict:
105    """
106    Serializes a Python collection into JSON by cases. See the README for the
107    precise list of supported types. The return dict has the following
108    structure:
109
110    - `collections.deque`:
111
112        ```py
113        {
114            "__type__": "collections.deque",
115            "__version__": 2,
116            "data": [...],
117            "maxlen": <int or None>,
118        }
119        ```
120
121    - `collections.namedtuple`
122
123        ```py
124        {
125            "__type__": "collections.namedtuple",
126            "__version__": 2,
127            "class": <str>,
128            "data": {...},
129        }
130        ```
131
132    - `set`
133
134        ```py
135        {
136            "__type__": "collections.set",
137            "__version__": 2,
138            "data": [...],
139        }
140        ```
141
142    - `tuple`
143
144        ```py
145        {
146            "__type__": "collections.tuple",
147            "__version__": 1,
148            "data": [...],
149        }
150        ```
151
152    """
153    encoders: list[Tuple[type, Callable[[Any, Context], dict]]] = [
154        (deque, _deque_to_json),
155        (tuple, _tuple_to_json),
156        (set, _set_to_json),
157    ]
158    for t, f in encoders:
159        if isinstance(obj, t):
160            return f(obj, ctx)
161    raise TypeNotSupported()