1# Copyright 2017-2026 Alexey Stepanov aka penguinolog
2
3# Licensed under the Apache License, Version 2.0 (the "License"); you may
4# not use this file except in compliance with the License. You may obtain
5# a copy of the License at
6
7# http://www.apache.org/licenses/LICENSE-2.0
8
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations
13# under the License.
14
15"""JSONField implementation for SQLAlchemy."""
16
17from __future__ import annotations
18
19import json
20import re
21import typing
22
23import sqlalchemy.ext.mutable
24import sqlalchemy.types
25
26if typing.TYPE_CHECKING:
27 import types
28 from collections.abc import Mapping
29
30 from sqlalchemy.engine import Dialect
31 from sqlalchemy.sql.type_api import TypeEngine
32
33__all__ = ("JSONField", "mutable_json_field")
34
35_NON_ASCII = re.compile(r"[^\x00-\x7f]")
36
37
38def _escape_non_ascii_char(match: re.Match[str]) -> str:
39 r"""Escape a single non-ASCII character as JSON ``\uXXXX`` sequence (surrogate pair outside of the BMP).
40
41 :param match: regex match of exactly one non-ASCII character
42 :type match: re.Match[str]
43 :return: escape sequence, identical to the standard ``json`` module output
44 :rtype: str
45 """
46 if (code := ord(match.group())) < 0x10000:
47 return f"\\u{code:04x}"
48 code -= 0x10000
49 return f"\\u{0xD800 | (code >> 10):04x}\\u{0xDC00 | (code & 0x3FF):04x}"
50
51
52def _escape_non_ascii(encoded: str) -> str:
53 """Make encoded JSON ASCII-only.
54
55 In valid JSON non-ASCII characters can appear only inside string literals,
56 where an escape sequence is equivalent to the character itself.
57
58 :param encoded: encoded JSON document
59 :type encoded: str
60 :return: the same document with all non-ASCII characters escaped
61 :rtype: str
62 """
63 if encoded.isascii():
64 return encoded
65 return _NON_ASCII.sub(_escape_non_ascii_char, encoded)
66
67
68_SURROGATE = re.compile(r"[\ud800-\udfff]")
69
70
71def _join_surrogates(value: typing.Any) -> typing.Any:
72 """Join UTF-16 surrogate pairs left by a decoder into the real characters in all decoded strings.
73
74 Unpaired surrogates are kept as is.
75
76 :param value: decoded JSON value
77 :type value: typing.Any
78 :return: the same structure with fixed strings
79 :rtype: typing.Any
80 """
81 if isinstance(value, str):
82 if _SURROGATE.search(value) is None:
83 return value
84 return value.encode("utf-16-le", "surrogatepass").decode("utf-16-le", "surrogatepass")
85 if isinstance(value, dict):
86 return {_join_surrogates(key): _join_surrogates(item) for key, item in value.items()}
87 if isinstance(value, list):
88 return [_join_surrogates(item) for item in value]
89 return value
90
91
92def _is_ujson5(codec: typing.Any) -> bool:
93 """Check whether JSON library is ujson5 (or its submodule) and requires decoder workarounds.
94
95 :param codec: JSON library
96 :type codec: typing.Any
97 :return: codec is ujson5
98 :rtype: bool
99 """
100 name = getattr(codec, "__name__", "")
101 return isinstance(name, str) and name.partition(".")[0] == "ujson5"
102
103
104# noinspection PyAbstractClass
105class JSONField(sqlalchemy.types.TypeDecorator): # type: ignore[type-arg] # pylint: disable=too-many-instance-attributes
106 r"""Represent an immutable structure as a json-encoded string or json.
107
108 Usage::
109
110 JSONField(enforce_string=True | False, enforce_unicode=True | False)
111
112 JSON library is any object with ``dumps`` and ``loads`` callables (module, class, namespace).
113 Libraries with a non-standard ``dumps`` signature are supported via extra arguments, for example ``orjson``::
114
115 JSONField(enforce_string=True, json=orjson, skip_ensure_ascii=True, decode_bytes=True)
116
117 ``ujson5`` is detected automatically and its decoder flaws are worked around:
118 ``strict=False`` is used by default (escaped ``\n``, ``\t`` are rejected in the strict mode)
119 and escaped surrogate pairs (characters outside the BMP, like emoji) are joined after decoding.
120 """
121
122 def process_literal_param(self, value: typing.Any, dialect: Dialect) -> typing.Any:
123 """Re-use of process_bind_param.
124
125 :return: encoded value if required
126 :rtype: typing.Union[str, typing.Any]
127 """
128 return self.process_bind_param(value, dialect)
129
130 impl = sqlalchemy.types.TypeEngine # Special placeholder
131 cache_ok = False # Cache complexity due to requirement of value re-serialization and mutability
132
133 def __init__( # pylint: disable=keyword-arg-before-vararg
134 self,
135 enforce_string: bool = False,
136 enforce_unicode: bool = False,
137 json: types.ModuleType | typing.Any = json, # pylint: disable=redefined-outer-name
138 json_type: TypeEngine[typing.Any] | type[TypeEngine[typing.Any]] = sqlalchemy.JSON,
139 *args: typing.Any,
140 skip_ensure_ascii: bool = False,
141 decode_bytes: bool = False,
142 dumps_kwargs: Mapping[str, typing.Any] | None = None,
143 loads_kwargs: Mapping[str, typing.Any] | None = None,
144 **kwargs: typing.Any,
145 ) -> None:
146 """JSONField.
147
148 :param enforce_string: enforce String(UnicodeText) type usage
149 :type enforce_string: bool
150 :param enforce_unicode: do not encode non-ascii data
151 :type enforce_unicode: bool
152 :param json: JSON encoding/decoding library. By default: standard json package.
153 :param json_type: the sqlalchemy/dialect class that will be used to render the DB JSON type.
154 By default: sqlalchemy.JSON
155 :param args: extra baseclass arguments
156 :type args: typing.Any
157 :param skip_ensure_ascii: do not pass ``ensure_ascii`` to ``dumps`` (library does not support it, like orjson).
158 Library is expected to produce UTF-8 output; if ``enforce_unicode`` is not set,
159 non-ASCII characters are escaped after encoding.
160 :type skip_ensure_ascii: bool
161 :param decode_bytes: ``dumps`` returns ``bytes`` (like orjson): decode result as UTF-8.
162 :type decode_bytes: bool
163 :param dumps_kwargs: extra keyword arguments for ``dumps`` (like ujson5 ``key_quotation`` or orjson ``option``).
164 ``ensure_ascii`` is not allowed: it is controlled by ``enforce_unicode``.
165 :type dumps_kwargs: typing.Mapping[str, typing.Any] | None
166 :param loads_kwargs: extra keyword arguments for ``loads``. ujson5: ``{"strict": False}`` is used by default.
167 :type loads_kwargs: typing.Mapping[str, typing.Any] | None
168 :param kwargs: extra baseclass keyworded arguments
169 :type kwargs: typing.Any
170 :raises ValueError: ``ensure_ascii`` is passed in ``dumps_kwargs``
171 """
172 dumps_kwargs = dict(dumps_kwargs or {})
173 if "ensure_ascii" in dumps_kwargs:
174 raise ValueError("ensure_ascii is controlled by enforce_unicode, do not pass it in dumps_kwargs")
175 if not skip_ensure_ascii:
176 dumps_kwargs["ensure_ascii"] = not enforce_unicode
177
178 self.__escape_non_ascii = skip_ensure_ascii and not enforce_unicode
179 self.__decode_bytes = decode_bytes
180 self.__dumps_kwargs = dumps_kwargs
181 self.__join_surrogates = _is_ujson5(json)
182 if self.__join_surrogates:
183 # ujson5 strict mode rejects valid escaped control characters (\n, \t, \r, \0) in strings
184 self.__loads_kwargs: dict[str, typing.Any] = {"strict": False, **(loads_kwargs or {})}
185 else:
186 self.__loads_kwargs = dict(loads_kwargs or {})
187 self.__enforce_string = enforce_string
188 self.__json_codec = json
189 self.__json_type = json_type
190 super().__init__(*args, **kwargs)
191
192 def __use_json(self, dialect: Dialect) -> bool:
193 """Determine which encoder to use.
194
195 :return: use engine-based json encoder
196 :rtype: bool
197 """
198 return hasattr(dialect, "_json_serializer") and not self.__enforce_string
199
200 def load_dialect_impl(self, dialect: Dialect) -> TypeEngine[typing.Any]:
201 """Select impl by dialect.
202
203 :return: dialect implementation depends on decoding method
204 :rtype: TypeEngine
205 """
206 # types are handled by DefaultDialect, Dialect class is abstract
207 if self.__use_json(dialect):
208 return dialect.type_descriptor(self.__json_type) # type: ignore[arg-type]
209 return dialect.type_descriptor(sqlalchemy.UnicodeText) # type: ignore[arg-type]
210
211 def process_bind_param(self, value: typing.Any, dialect: Dialect) -> str | typing.Any:
212 """Encode data, if required.
213
214 :return: encoded value if required
215 :rtype: typing.Union[str, typing.Any]
216 """
217 if self.__use_json(dialect) or value is None:
218 return value
219
220 encoded = self.__json_codec.dumps(value, **self.__dumps_kwargs)
221 if self.__decode_bytes and isinstance(encoded, (bytes, bytearray, memoryview)):
222 encoded = bytes(encoded).decode("utf-8")
223 if self.__escape_non_ascii:
224 return _escape_non_ascii(encoded)
225 return encoded
226
227 def process_result_value(self, value: str | typing.Any, dialect: Dialect) -> typing.Any:
228 """Decode data, if required.
229
230 :return: decoded result value if required
231 :rtype: typing.Any
232 """
233 if self.__use_json(dialect) or value is None:
234 return value
235
236 decoded = self.__json_codec.loads(value, **self.__loads_kwargs)
237 if self.__join_surrogates:
238 return _join_surrogates(decoded)
239 return decoded
240
241
242def mutable_json_field( # pylint: disable=keyword-arg-before-vararg, redefined-outer-name
243 enforce_string: bool = False,
244 enforce_unicode: bool = False,
245 json: types.ModuleType | typing.Any = json,
246 *args: typing.Any,
247 skip_ensure_ascii: bool = False,
248 decode_bytes: bool = False,
249 dumps_kwargs: Mapping[str, typing.Any] | None = None,
250 loads_kwargs: Mapping[str, typing.Any] | None = None,
251 **kwargs: typing.Any,
252) -> JSONField:
253 """Mutable JSONField creator.
254
255 :param enforce_string: enforce String(UnicodeText) type usage
256 :type enforce_string: bool
257 :param enforce_unicode: do not encode non-ascii data
258 :type enforce_unicode: bool
259 :param json: JSON encoding/decoding library.
260 By default: standard json package.
261 :param args: extra baseclass arguments
262 :type args: typing.Any
263 :param skip_ensure_ascii: do not pass ``ensure_ascii`` to ``dumps`` (library does not support it, like orjson).
264 :type skip_ensure_ascii: bool
265 :param decode_bytes: ``dumps`` returns ``bytes`` (like orjson): decode result as UTF-8.
266 :type decode_bytes: bool
267 :param dumps_kwargs: extra keyword arguments for ``dumps``.
268 :type dumps_kwargs: typing.Mapping[str, typing.Any] | None
269 :param loads_kwargs: extra keyword arguments for ``loads``. ujson5: ``{"strict": False}`` is used by default.
270 :type loads_kwargs: typing.Mapping[str, typing.Any] | None
271 :param kwargs: extra baseclass keyworded arguments
272 :type kwargs: typing.Any
273 :return: Mutable JSONField via MutableDict.as_mutable
274 :rtype: JSONField
275 """
276 return sqlalchemy.ext.mutable.MutableDict.as_mutable( # type: ignore[return-value]
277 JSONField( # type: ignore[misc]
278 enforce_string=enforce_string,
279 enforce_unicode=enforce_unicode,
280 json=json,
281 *args, # noqa: B026
282 skip_ensure_ascii=skip_ensure_ascii,
283 decode_bytes=decode_bytes,
284 dumps_kwargs=dumps_kwargs,
285 loads_kwargs=loads_kwargs,
286 **kwargs,
287 )
288 )