/
/
/
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 serialize_to_json(obj: Any) -> Any:
49 """Serialize a value (or a list of values) to json."""
50 if obj is None:
51 return obj
52 if hasattr(obj, "to_json"):
53 return obj.to_json()
54 return json_dumps(get_serializable_value(obj))
55
56
57def json_dumps(data: Any, indent: bool = False) -> str:
58 """Dump json string."""
59 # we use the passthrough dataclass option because we use mashumaro for that
60 option = orjson.OPT_OMIT_MICROSECONDS | orjson.OPT_PASSTHROUGH_DATACLASS
61 if indent:
62 option |= orjson.OPT_INDENT_2
63 return orjson.dumps(
64 data,
65 default=_json_default,
66 option=option,
67 ).decode("utf-8")
68
69
70async def async_json_dumps(data: Any, indent: bool = False) -> str:
71 """Dump json string async."""
72 return await asyncio.to_thread(json_dumps, data, indent)
73
74
75json_loads = orjson.loads
76
77
78async def async_json_loads(data: str) -> Any:
79 """Load json string async."""
80 return await asyncio.to_thread(json_loads, data)
81
82
83def strip_code_fence(text: str) -> str:
84 """
85 Remove a markdown code fence that wraps an entire text, plus surrounding whitespace.
86
87 Anything else is returned with only its whitespace trimmed: a fence around part of
88 the text, more than one fenced block, an unterminated fence, a fence whose opening
89 line carries more than a language tag, or a body containing the fence itself. Text
90 that is not exactly one fenced block therefore stays as invalid to a strict parser
91 as it was before.
92
93 :param text: Text that may be wrapped in a markdown code fence.
94 :return: The fenced content, or the trimmed text when it is not a single fenced block.
95 """
96 stripped = text.strip()
97 fence_length = len(stripped) - len(stripped.lstrip("`"))
98 if fence_length < _MIN_CODE_FENCE_LENGTH:
99 return stripped
100 fence = "`" * fence_length
101 inner = stripped[fence_length:]
102 if not inner.endswith(fence):
103 return stripped
104 inner = inner[:-fence_length]
105 # A closing fence must be the final run of backticks, matching the opening length.
106 if inner.endswith("`"):
107 return stripped
108 tag, separator, body = inner.partition("\n")
109 if not separator or not _CODE_FENCE_TAG_PATTERN.fullmatch(tag.strip()):
110 return stripped
111 if fence in body:
112 return stripped
113 return body.strip()
114
115
116TargetT = TypeVar("TargetT", bound=DataClassORJSONMixin)
117
118
119async def load_json_file[TargetT: DataClassORJSONMixin](
120 path: str, target_class: type[TargetT]
121) -> TargetT:
122 """Load JSON from file."""
123 async with aiofiles.open(path) as _file:
124 content = await _file.read()
125 return target_class.from_json(content)
126
127
128async def load_json_dict(path: str) -> dict[str, Any]:
129 """
130 Load a plain JSON object from file as a dict.
131
132 For JSON files that are not backed by a dataclass (e.g. translation strings files).
133
134 :param path: Absolute path to the JSON file.
135 """
136 async with aiofiles.open(path, "rb") as _file:
137 content = await _file.read()
138 data = orjson.loads(content)
139 if not isinstance(data, dict):
140 msg = f"Expected a JSON object in {path}, got {type(data).__name__}"
141 raise TypeError(msg)
142 return data
143
144
145def _json_default(obj: Any) -> Any:
146 """Convert a value for orjson, raising a descriptive error for unhandled types."""
147 value = get_serializable_value(obj)
148 if value is obj:
149 cls = type(obj)
150 msg = (
151 f"unhandled type for json serialization: {cls.__module__}.{cls.__qualname__}"
152 " - pass a dict (e.g. via .to_dict()) instead of the raw object"
153 )
154 raise TypeError(msg)
155 return value
156