Coverage for open_webui/utils/json_codec.py: 19%

1"""The app-wide JSON codec, selected by the ``ENABLE_ORJSON`` env var. 

2 

3Every module that would otherwise reach for stdlib ``json`` imports ``JSONCodec`` 

4from here, so the whole app switches implementation from a single flag. With the 

5flag off these are stdlib ``json`` and engineio's codec verbatim, so the default 

6behaviour is exactly what it was before orjson entered the picture. ``dumps_bytes`` 

7returns UTF-8 bytes for sinks that re-parse the payload; under orjson it skips 

8both the str round trip and the line-separator escaping ``dumps`` applies, so 

9never feed it to line-framed output such as SSE. 

10""" 

11 

12from __future__ import annotations 

13 

14import json as stdlib_json 

15 

16from engineio import json as engineio_json 

17from open_webui.env import ENABLE_ORJSON 

18 

19if ENABLE_ORJSON: 19 ↛ 20line 19 didn't jump to line 20 because the condition on line 19 was never true

20 import orjson 

21 

22 # Module-level because CPython rebuilds these dicts on every call. 

23 FAST_PATH_KWARGS = ({'separators': (',', ':')}, {'ensure_ascii': False}) 

24 

25 class ORJSONCodec: 

26 """stdlib-``json``-compatible codec backed by orjson. 

27 

28 The fast path is not byte-for-byte stdlib: it is always compact, formats 

29 floats orjson's way (``1e16``, not ``1e+16``), and is raw UTF-8 apart from 

30 the three line separators ``dumps`` escapes, so a ``separators`` caller loses 

31 stdlib's ASCII escaping and an ``ensure_ascii=False`` caller loses its 

32 spacing. ``dumps`` also serializes ``datetime``/``UUID``/dataclasses that 

33 stdlib refuses, and encodes ``NaN``/``Infinity`` as ``null``. ``loads`` 

34 decodes integers above ``2**64-1`` or below ``-2**63`` as ``float`` and does 

35 not enforce engineio's 100-digit integer-literal limit. 

36 

37 What orjson does reject (non-str dict keys and oversized ints on ``dumps``, 

38 the ``NaN``/``Infinity`` literals on ``loads``) falls back to engineio's 

39 stdlib-based codec, and with it stdlib's formatting. 

40 """ 

41 

42 JSONDecodeError = engineio_json.JSONDecodeError 

43 

44 @staticmethod 

45 def dumps(obj, *args, **kwargs): 

46 if args or (kwargs and kwargs not in FAST_PATH_KWARGS): 

47 return engineio_json.dumps(obj, *args, **kwargs) 

48 try: 

49 serialized = orjson.dumps(obj).decode('utf-8') 

50 except (TypeError, ValueError): 

51 return engineio_json.dumps(obj, *args, **kwargs) 

52 # Raw, these three split an SSE frame reassembled with ``splitlines()``. 

53 # A dict-table translate walks char by char; chained replace runs on C fast paths. 

54 if '\u2028' in serialized or '\u2029' in serialized or '\x85' in serialized: 

55 return serialized.replace('\u2028', '\\u2028').replace('\u2029', '\\u2029').replace('\x85', '\\u0085') 

56 return serialized 

57 

58 @staticmethod 

59 def loads(s, *args, **kwargs): 

60 if args or kwargs: 

61 return engineio_json.loads(s, *args, **kwargs) 

62 try: 

63 return orjson.loads(s) 

64 except (TypeError, ValueError): 

65 return engineio_json.loads(s, *args, **kwargs) 

66 

67 # Drop-in for stdlib ``json``: ``JSONCodec.dumps`` / ``JSONCodec.loads``. 

68 JSONCodec = ORJSONCodec 

69 # Codec handed to the socket.io/engineio managers, which default to their own. 

70 SOCKETIO_JSON = ORJSONCodec 

71 

72 def dumps_bytes(obj) -> bytes: 

73 """JSON as UTF-8 bytes, skipping the str round trip and the escaping ``dumps`` does.""" 

74 try: 

75 return orjson.dumps(obj) 

76 except (TypeError, ValueError): 

77 return engineio_json.dumps(obj).encode('utf-8') 

78else: 

79 JSONCodec = stdlib_json 

80 SOCKETIO_JSON = engineio_json 

81 

82 def dumps_bytes(obj) -> bytes: 

83 """JSON as UTF-8 bytes; here simply ``dumps`` encoded.""" 

84 return stdlib_json.dumps(obj).encode('utf-8')