-
Notifications
You must be signed in to change notification settings - Fork 238
Expand file tree
/
Copy pathjson.py
More file actions
99 lines (79 loc) · 3.12 KB
/
json.py
File metadata and controls
99 lines (79 loc) · 3.12 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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import json
from contextlib import nullcontext
from typing import Union, TextIO, TYPE_CHECKING, Type, List
if TYPE_CHECKING: # pragma: no cover
from docarray.typing import T
class JsonIOMixin:
"""Save/load a array into a JSON file."""
def save_json(
self,
file: Union[str, TextIO],
protocol: str = 'jsonschema',
encoding: str = 'utf-8',
**kwargs
) -> None:
"""Save array elements into a JSON file.
Comparing to :meth:`save_binary`, it is human-readable but slower to save/load and the file size larger.
:param file: File or filename to which the data is saved.
:param protocol: `jsonschema` or `protobuf`
:param encoding: encoding used to save data into a JSON file. By default, ``utf-8`` is used.
"""
if hasattr(file, 'write'):
file_ctx = nullcontext(file)
else:
file_ctx = open(file, 'w', encoding=encoding)
with file_ctx as fp:
fp.write(self.to_json(protocol=protocol, **kwargs))
@classmethod
def load_json(
cls: Type['T'],
file: Union[str, TextIO],
protocol: str = 'jsonschema',
encoding: str = 'utf-8',
**kwargs
) -> 'T':
"""Load array elements from a JSON file.
:param file: File or filename or a JSON string to which the data is saved.
:param protocol: `jsonschema` or `protobuf`
:param encoding: encoding used to load data from a JSON file. By default, ``utf-8`` is used.
:return: a DocumentArrayLike object
"""
if hasattr(file, 'read'):
file_ctx = nullcontext(file)
else:
file_ctx = open(file, 'r', encoding=encoding)
with file_ctx as fp:
return cls.from_json(fp.read(), protocol=protocol, **kwargs)
@classmethod
def from_json(
cls: Type['T'],
file: Union[str, bytes, bytearray],
protocol: str = 'jsonschema',
**kwargs
) -> 'T':
from docarray import Document
json_docs = json.loads(file)
return cls(
[Document.from_dict(v, protocol=protocol) for v in json_docs], **kwargs
)
@classmethod
def from_list(
cls: Type['T'], values: List, protocol: str = 'jsonschema', **kwargs
) -> 'T':
from docarray import Document
return cls(Document.from_dict(v, protocol=protocol, **kwargs) for v in values)
def to_list(self, protocol: str = 'jsonschema', **kwargs) -> List:
"""Convert the object into a Python list.
:param protocol: `jsonschema` or `protobuf`
:return: a Python list
"""
return [d.to_dict(protocol=protocol, **kwargs) for d in self]
def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str:
"""Convert the object into a JSON string. Can be loaded via :meth:`.load_json`.
:param protocol: `jsonschema` or `protobuf`
:return: a Python list
"""
return json.dumps(self.to_list(protocol=protocol, **kwargs))
# to comply with Document interfaces but less semantically accurate
to_dict = to_list
from_dict = from_list