/
/
/
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) -> Any:
30 """Parse the value to its serializable equivalent."""
31 if getattr(obj, "do_not_serialize", None):
32 return None
33 if isinstance(obj, list | set | filter | tuple | dict_values | dict_keys) or (
34 obj.__class__.__name__ == "dict_valueiterator"
35 ):
36 return [get_serializable_value(x) for x in obj]
37 if hasattr(obj, "to_dict"):
38 return obj.to_dict()
39 if isinstance(obj, bytes):
40 return base64.b64encode(obj).decode("ascii")
41 if isinstance(obj, DO_NOT_SERIALIZE_TYPES):
42 return None
43 # unhandled values are returned as-is on purpose: serialize_to_json and the
44 # recursion above rely on natively serializable values passing through
45 return obj
46
47
48def make_utf8_safe(value: str) -> str:
49 """
50 Return the given string made encodable if it is not valid UTF-8.
51
52 Lone surrogates, as the os module returns for undecodable filesystem paths, come
53 back as their backslash escapes.
54
55 :param value: The string to make safe.
56 """
57 try:
58 value.encode()
59 except UnicodeEncodeError:
60 try:
61 # a surrogate escape stands for one raw byte, so map it back to show
62 # the actual filename bytes instead of the surrogates
63 raw = value.encode("utf-8", "surrogateescape")
64 except UnicodeEncodeError:
65 # any other lone surrogate has no byte to map back to
66 return value.encode("utf-8", "backslashreplace").decode()
67 return raw.decode("utf-8", "backslashreplace")
68 return value
69
70
71def serialize_to_json(obj: Any) -> Any:
72 """Serialize a value (or a list of values) to json."""
73 if obj is None:
74 return obj
75 if hasattr(obj, "to_json"):
76 return obj.to_json()
77 return json_dumps(get_serializable_value(obj))
78
79
80def json_dumps(data: Any, indent: bool = False) -> str:
81 """Dump json string."""
82 # we use the passthrough dataclass option because we use mashumaro for that
83 option = orjson.OPT_OMIT_MICROSECONDS | orjson.OPT_PASSTHROUGH_DATACLASS
84 if indent:
85 option |= orjson.OPT_INDENT_2
86 return orjson.dumps(
87 data,
88 default=_json_default,
89 option=option,
90 ).decode("utf-8")
91
92
93async def async_json_dumps(data: Any, indent: bool = False) -> str:
94 """Dump json string async."""
95 return await asyncio.to_thread(json_dumps, data, indent)
96
97
98json_loads = orjson.loads
99
100
101async def async_json_loads(data: str) -> Any:
102 """Load json string async."""
103 return await asyncio.to_thread(json_loads, data)
104
105
106def strip_code_fence(text: str) -> str:
107 """
108 Remove a markdown code fence that wraps an entire text, plus surrounding whitespace.
109
110 Anything else is returned with only its whitespace trimmed: a fence around part of
111 the text, more than one fenced block, an unterminated fence, a fence whose opening
112 line carries more than a language tag, or a body containing the fence itself. Text
113 that is not exactly one fenced block therefore stays as invalid to a strict parser
114 as it was before.
115
116 :param text: Text that may be wrapped in a markdown code fence.
117 :return: The fenced content, or the trimmed text when it is not a single fenced block.
118 """
119 stripped = text.strip()
120 fence_length = len(stripped) - len(stripped.lstrip("`"))
121 if fence_length < _MIN_CODE_FENCE_LENGTH:
122 return stripped
123 fence = "`" * fence_length
124 inner = stripped[fence_length:]
125 if not inner.endswith(fence):
126 return stripped
127 inner = inner[:-fence_length]
128 # A closing fence must be the final run of backticks, matching the opening length.
129 if inner.endswith("`"):
130 return stripped
131 tag, separator, body = inner.partition("\n")
132 if not separator or not _CODE_FENCE_TAG_PATTERN.fullmatch(tag.strip()):
133 return stripped
134 if fence in body:
135 return stripped
136 return body.strip()
137
138
139TargetT = TypeVar("TargetT", bound=DataClassORJSONMixin)
140
141
142async def load_json_file[TargetT: DataClassORJSONMixin](
143 path: str, target_class: type[TargetT]
144) -> TargetT:
145 """Load JSON from file."""
146 async with aiofiles.open(path) as _file:
147 content = await _file.read()
148 return target_class.from_json(content)
149
150
151async def load_json_dict(path: str) -> dict[str, Any]:
152 """
153 Load a plain JSON object from file as a dict.
154
155 For JSON files that are not backed by a dataclass (e.g. translation strings files).
156
157 :param path: Absolute path to the JSON file.
158 """
159 async with aiofiles.open(path, "rb") as _file:
160 content = await _file.read()
161 data = orjson.loads(content)
162 if not isinstance(data, dict):
163 msg = f"Expected a JSON object in {path}, got {type(data).__name__}"
164 raise TypeError(msg)
165 return data
166
167
168def _json_default(obj: Any) -> Any:
169 """Convert a value for orjson, raising a descriptive error for unhandled types."""
170 value = get_serializable_value(obj)
171 if value is obj:
172 cls = type(obj)
173 msg = (
174 f"unhandled type for json serialization: {cls.__module__}.{cls.__qualname__}"
175 " - pass a dict (e.g. via .to_dict()) instead of the raw object"
176 )
177 raise TypeError(msg)
178 return value
179