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 turbo_broccoli.context import Context
  7from turbo_broccoli.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
 90# pylint: disable=missing-function-docstring
 91def from_json(dct: dict, ctx: Context) -> Any:
 92    decoders = {
 93        "collections.deque": _json_to_deque,
 94        "collections.namedtuple": _json_to_namedtuple,
 95        "collections.set": _json_to_set,
 96        "collections.tuple": _json_to_tuple,
 97    }
 98    try:
 99        type_name = dct["__type__"]
100        return decoders[type_name](dct, ctx)
101    except KeyError as exc:
102        raise DeserializationError() from exc
103
104
105def to_json(obj: Any, ctx: Context) -> dict:
106    """
107    Serializes a Python collection into JSON by cases. See the README for the
108    precise list of supported types. The return dict has the following
109    structure:
110
111    - `collections.deque`:
112
113        ```py
114        {
115            "__type__": "collections.deque",
116            "__version__": 2,
117            "data": [...],
118            "maxlen": <int or None>,
119        }
120        ```
121
122    - `collections.namedtuple`
123
124        ```py
125        {
126            "__type__": "collections.namedtuple",
127            "__version__": 2,
128            "class": <str>,
129            "data": {...},
130        }
131        ```
132
133    - `set`
134
135        ```py
136        {
137            "__type__": "collections.set",
138            "__version__": 2,
139            "data": [...],
140        }
141        ```
142
143    - `tuple`
144
145        ```py
146        {
147            "__type__": "collections.tuple",
148            "__version__": 1,
149            "data": [...],
150        }
151        ```
152
153    """
154    encoders: list[Tuple[type, Callable[[Any, Context], dict]]] = [
155        (deque, _deque_to_json),
156        (tuple, _tuple_to_json),
157        (set, _set_to_json),
158    ]
159    for t, f in encoders:
160        if isinstance(obj, t):
161            return f(obj, ctx)
162    raise TypeNotSupported()