Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy_jsonfield/jsonfield.py: 52%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

84 statements  

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 )