/
/
/
1"""Helpers to work with (de)serializing of json."""
2
3import asyncio
4import base64
5import re
6from _collections_abc import dict_keys, dict_values
7from types import MethodType
8from typing import Any, TypeVar
9
10import aiofiles
11import orjson
12from mashumaro.mixins.orjson import DataClassORJSONMixin
13
14JSON_ENCODE_EXCEPTIONS = (TypeError, ValueError)
15JSON_DECODE_EXCEPTIONS = (orjson.JSONDecodeError,)
16
17DO_NOT_SERIALIZE_TYPES = (MethodType, asyncio.Task)
18
19_MIN_CODE_FENCE_LENGTH = 3
20# Applied to the opening line only, so a fence that carries anything other than a bare
21# language tag is left alone instead of having that line silently discarded.
22_CODE_FENCE_TAG_PATTERN = re.compile(r"[A-Za-z0-9+#._-]*")
23
24# Type alias for plain, JSON-serializable data.
25# Note: tuples are accepted but will be returned as lists after a JSON round-trip.
26SerializableType = str | int | float | bool | None | list[Any] | tuple[Any, ...] | dict[str, Any]
27
28
29def get_serializable_value(obj: Any, raise_unhandled: bool = False) -> Any:
30 """Parse the value to its serializable equivalent."""
31 if getattr(obj, "do_not_serialize", None):
32 return None
33 if (
34 isinstance(obj, list | set | filter | tuple | dict_values | dict_keys | dict_values)
35 or obj.__class__ == "dict_valueiterator"
36 ):
37 return [get_serializable_value(x) for x in obj]
38 if hasattr(obj, "to_dict"):
39 return obj.to_dict()
40 if isinstance(obj, bytes):
41 return base64.b64encode(obj).decode("ascii")
42 if isinstance(obj, DO_NOT_SERIALIZE_TYPES):
43 return None
44 if raise_unhandled:
45 raise TypeError
46 return obj
47
48
49def serialize_to_json(obj: Any) -> Any:
50 """Serialize a value (or a list of values) to json."""
51 if obj is None:
52 return obj
53 if hasattr(obj, "to_json"):
54 return obj.to_json()
55 return json_dumps(get_serializable_value(obj))
56
57
58def json_dumps(data: Any, indent: bool = False) -> str:
59 """Dump json string."""
60 # we use the passthrough dataclass option because we use mashumaro for that
61 option = orjson.OPT_OMIT_MICROSECONDS | orjson.OPT_PASSTHROUGH_DATACLASS
62 if indent:
63 option |= orjson.OPT_INDENT_2
64 return orjson.dumps(
65 data,
66 default=get_serializable_value,
67 option=option,
68 ).decode("utf-8")
69
70
71async def async_json_dumps(data: Any, indent: bool = False) -> str:
72 """Dump json string async."""
73 return await asyncio.to_thread(json_dumps, data, indent)
74
75
76json_loads = orjson.loads
77
78
79async def async_json_loads(data: str) -> Any:
80 """Load json string async."""
81 return await asyncio.to_thread(json_loads, data)
82
83
84def strip_code_fence(text: str) -> str:
85 """
86 Remove a markdown code fence that wraps an entire text, plus surrounding whitespace.
87
88 Anything else is returned with only its whitespace trimmed: a fence around part of
89 the text, more than one fenced block, an unterminated fence, a fence whose opening
90 line carries more than a language tag, or a body containing the fence itself. Text
91 that is not exactly one fenced block therefore stays as invalid to a strict parser
92 as it was before.
93
94 :param text: Text that may be wrapped in a markdown code fence.
95 :return: The fenced content, or the trimmed text when it is not a single fenced block.
96 """
97 stripped = text.strip()
98 fence_length = len(stripped) - len(stripped.lstrip("`"))
99 if fence_length < _MIN_CODE_FENCE_LENGTH:
100 return stripped
101 fence = "`" * fence_length
102 inner = stripped[fence_length:]
103 if not inner.endswith(fence):
104 return stripped
105 inner = inner[:-fence_length]
106 # A closing fence must be the final run of backticks, matching the opening length.
107 if inner.endswith("`"):
108 return stripped
109 tag, separator, body = inner.partition("\n")
110 if not separator or not _CODE_FENCE_TAG_PATTERN.fullmatch(tag.strip()):
111 return stripped
112 if fence in body:
113 return stripped
114 return body.strip()
115
116
117TargetT = TypeVar("TargetT", bound=DataClassORJSONMixin)
118
119
120async def load_json_file[TargetT: DataClassORJSONMixin](
121 path: str, target_class: type[TargetT]
122) -> TargetT:
123 """Load JSON from file."""
124 async with aiofiles.open(path) as _file:
125 content = await _file.read()
126 return target_class.from_json(content)
127
128
129async def load_json_dict(path: str) -> dict[str, Any]:
130 """
131 Load a plain JSON object from file as a dict.
132
133 For JSON files that are not backed by a dataclass (e.g. translation strings files).
134
135 :param path: Absolute path to the JSON file.
136 """
137 async with aiofiles.open(path, "rb") as _file:
138 content = await _file.read()
139 data = orjson.loads(content)
140 if not isinstance(data, dict):
141 msg = f"Expected a JSON object in {path}, got {type(data).__name__}"
142 raise TypeError(msg)
143 return data
144