/
/
/
1"""Helpers to work with (de)serializing of json."""
2
3import asyncio
4import base64
5from _collections_abc import dict_keys, dict_values
6from types import MethodType
7from typing import Any, TypeVar
8
9import aiofiles
10import orjson
11from mashumaro.mixins.orjson import DataClassORJSONMixin
12
13JSON_ENCODE_EXCEPTIONS = (TypeError, ValueError)
14JSON_DECODE_EXCEPTIONS = (orjson.JSONDecodeError,)
15
16DO_NOT_SERIALIZE_TYPES = (MethodType, asyncio.Task)
17
18
19def get_serializable_value(obj: Any, raise_unhandled: bool = False) -> Any:
20 """Parse the value to its serializable equivalent."""
21 if getattr(obj, "do_not_serialize", None):
22 return None
23 if (
24 isinstance(obj, list | set | filter | tuple | dict_values | dict_keys | dict_values)
25 or obj.__class__ == "dict_valueiterator"
26 ):
27 return [get_serializable_value(x) for x in obj]
28 if hasattr(obj, "to_dict"):
29 return obj.to_dict()
30 if isinstance(obj, bytes):
31 return base64.b64encode(obj).decode("ascii")
32 if isinstance(obj, DO_NOT_SERIALIZE_TYPES):
33 return None
34 if raise_unhandled:
35 raise TypeError
36 return obj
37
38
39def serialize_to_json(obj: Any) -> Any:
40 """Serialize a value (or a list of values) to json."""
41 if obj is None:
42 return obj
43 if hasattr(obj, "to_json"):
44 return obj.to_json()
45 return json_dumps(get_serializable_value(obj))
46
47
48def json_dumps(data: Any, indent: bool = False) -> str:
49 """Dump json string."""
50 # we use the passthrough dataclass option because we use mashumaro for that
51 option = orjson.OPT_OMIT_MICROSECONDS | orjson.OPT_PASSTHROUGH_DATACLASS
52 if indent:
53 option |= orjson.OPT_INDENT_2
54 return orjson.dumps(
55 data,
56 default=get_serializable_value,
57 option=option,
58 ).decode("utf-8")
59
60
61async def async_json_dumps(data: Any, indent: bool = False) -> str:
62 """Dump json string async."""
63 return await asyncio.to_thread(json_dumps, data, indent)
64
65
66json_loads = orjson.loads
67
68
69async def async_json_loads(data: str) -> Any:
70 """Load json string async."""
71 return await asyncio.to_thread(json_loads, data)
72
73
74TargetT = TypeVar("TargetT", bound=DataClassORJSONMixin)
75
76
77async def load_json_file[TargetT: DataClassORJSONMixin](
78 path: str, target_class: type[TargetT]
79) -> TargetT:
80 """Load JSON from file."""
81 async with aiofiles.open(path) as _file:
82 content = await _file.read()
83 return target_class.from_json(content)
84