-
Notifications
You must be signed in to change notification settings - Fork 238
Expand file tree
/
Copy pathjson.py
More file actions
41 lines (31 loc) · 1.07 KB
/
json.py
File metadata and controls
41 lines (31 loc) · 1.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from typing import Any, Callable, Dict, Type
import orjson
from docarray.utils._internal.pydantic import is_pydantic_v2
if not is_pydantic_v2:
from pydantic.json import ENCODERS_BY_TYPE
else:
ENCODERS_BY_TYPE: Dict[Type[Any], Callable[[Any], Any]] = {
bytes: lambda o: o.decode(),
frozenset: list,
set: list,
}
def _default_orjson(obj):
"""
default option for orjson dumps.
:param obj:
:return: return a json compatible object
"""
from docarray.base_doc import BaseNode
if isinstance(obj, BaseNode):
return obj._docarray_to_json_compatible()
else:
for cls_, encoder in ENCODERS_BY_TYPE.items():
if isinstance(obj, cls_):
return encoder(obj)
return obj
def orjson_dumps(v, *, default=None) -> bytes:
# dumps to bytes using orjson
return orjson.dumps(v, default=_default_orjson, option=orjson.OPT_SERIALIZE_NUMPY)
def orjson_dumps_and_decode(v, *, default=None) -> str:
# dumps to str using orjson
return orjson_dumps(v, default=default).decode()