Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit a96787e

Browse files
hanxiaohanxiao
authored
feat: use schema-ed json as default to_dict/json (#57)
1 parent 86b03f3 commit a96787e

13 files changed

Lines changed: 212 additions & 85 deletions

File tree

‎docarray/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
__version__ = '0.2.1'
1+
__version__ = '0.3.0'
22

33
from .document import Document
44
from .array import DocumentArray

‎docarray/array/mixins/io/csv.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -72,15 +72,12 @@ def save_csv(
7272
if with_header:
7373
writer.writeheader()
7474

75-
from .... import Document
76-
7775
for d in self:
78-
_d = d
79-
if exclude_fields:
80-
_d = Document(d, copy=True)
81-
_d.pop(*exclude_fields)
82-
83-
pd = _d.to_dict()
76+
pd = d.to_dict(
77+
protocol='jsonschema',
78+
exclude=set(exclude_fields) if exclude_fields else None,
79+
exclude_none=True,
80+
)
8481
if flatten_tags:
8582
t = pd.pop('tags')
8683
pd.update({f'tag__{k}': v for k, v in t.items()})

‎docarray/array/mixins/io/json.py‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,15 @@
1010
class JsonIOMixin:
1111
"""Save/load a array into a JSON file."""
1212

13-
def save_json(self, file: Union[str, TextIO]) -> None:
13+
def save_json(
14+
self, file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs
15+
) -> None:
1416
"""Save array elements into a JSON file.
1517
1618
Comparing to :meth:`save_binary`, it is human-readable but slower to save/load and the file size larger.
1719
1820
:param file: File or filename to which the data is saved.
21+
:param protocol: `jsonschema` or `protobuf`
1922
"""
2023
if hasattr(file, 'write'):
2124
file_ctx = nullcontext(file)
@@ -24,15 +27,17 @@ def save_json(self, file: Union[str, TextIO]) -> None:
2427

2528
with file_ctx as fp:
2629
for d in self:
27-
json.dump(d.to_dict(), fp)
30+
json.dump(d.to_dict(protocol=protocol, **kwargs), fp)
2831
fp.write('\n')
2932

3033
@classmethod
31-
def load_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T':
34+
def load_json(
35+
cls: Type['T'], file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs
36+
) -> 'T':
3237
"""Load array elements from a JSON file.
3338
3439
:param file: File or filename or a JSON string to which the data is saved.
35-
40+
:param protocol: `jsonschema` or `protobuf`
3641
:return: a DocumentArrayLike object
3742
"""
3843

@@ -48,31 +53,38 @@ def load_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T':
4853
constructor = Document.from_dict
4954

5055
with file_ctx as fp:
51-
return cls(constructor(v) for v in fp)
56+
return cls(constructor(v, protocol=protocol, **kwargs) for v in fp)
5257

5358
@classmethod
54-
def from_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T':
55-
return cls.load_json(file)
59+
def from_json(
60+
cls: Type['T'], file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs
61+
) -> 'T':
62+
return cls.load_json(file, protocol=protocol, **kwargs)
5663

5764
@classmethod
58-
def from_list(cls: Type['T'], values: List) -> 'T':
65+
def from_list(
66+
cls: Type['T'], values: List, protocol: str = 'jsonschema', **kwargs
67+
) -> 'T':
5968
from .... import Document
6069

61-
return cls(Document.from_dict(v) for v in values)
70+
return cls(Document.from_dict(v, protocol=protocol, **kwargs) for v in values)
6271

63-
def to_list(self, strict: bool = True) -> List:
72+
def to_list(self, protocol: str = 'jsonschema', **kwargs) -> List:
6473
"""Convert the object into a Python list.
6574
66-
.. note::
67-
Array like object such as :class:`numpy.ndarray` will be converted to Python list.
68-
75+
:param protocol: `jsonschema` or `protobuf`
6976
:return: a Python list
7077
"""
71-
return [d.to_dict(strict=strict) for d in self]
78+
return [d.to_dict(protocol=protocol, **kwargs) for d in self]
7279

73-
def to_json(self) -> str:
80+
def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str:
7481
"""Convert the object into a JSON string. Can be loaded via :meth:`.load_json`.
7582
83+
:param protocol: `jsonschema` or `protobuf`
7684
:return: a Python list
7785
"""
78-
return json.dumps(self.to_list())
86+
return json.dumps(self.to_list(protocol=protocol, **kwargs))
87+
88+
# to comply with Document interfaces but less semantically accurate
89+
to_dict = to_list
90+
from_dict = from_list

‎docarray/base.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ def __init__(
5252
if _unknown_kwargs and unknown_fields_handler == 'catch':
5353
getattr(self, self._unresolved_fields_dest).update(_unknown_kwargs)
5454

55-
if _obj is None and not kwargs and self._data is None:
55+
if not _obj and not kwargs and self._data is None:
5656
self._data = self._data_class(self)
5757

5858
if self._data is None:

‎docarray/document/mixins/porting.py‎

Lines changed: 76 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import dataclasses
22
import pickle
3+
import warnings
34
from typing import Optional, TYPE_CHECKING, Type, Dict, Any
45
import base64
56

@@ -11,32 +12,76 @@
1112

1213
class PortingMixin:
1314
@classmethod
14-
def from_dict(cls: Type['T'], obj: Dict) -> 'T':
15-
from google.protobuf import json_format
16-
from ...proto.docarray_pb2 import DocumentProto
15+
def from_dict(
16+
cls: Type['T'], obj: Dict, protocol: str = 'jsonschema', **kwargs
17+
) -> 'T':
18+
"""Convert a dict object into a Document.
19+
20+
:param obj: a Python dict object
21+
:param protocol: `jsonschema` or `protobuf`
22+
:param kwargs: extra key-value args pass to pydantic and protobuf parser.
23+
:return: the parsed Document object
24+
"""
25+
if protocol == 'jsonschema':
26+
from ..pydantic_model import PydanticDocument
27+
28+
return cls.from_pydantic_model(PydanticDocument.parse_obj(obj, **kwargs))
29+
elif protocol == 'protobuf':
30+
from google.protobuf import json_format
31+
from ...proto.docarray_pb2 import DocumentProto
1732

18-
pb_msg = DocumentProto()
19-
json_format.ParseDict(obj, pb_msg)
20-
return cls.from_protobuf(pb_msg)
33+
pb_msg = DocumentProto()
34+
json_format.ParseDict(obj, pb_msg, **kwargs)
35+
return cls.from_protobuf(pb_msg)
36+
else:
37+
raise ValueError(f'protocol=`{protocol}` is not supported')
2138

2239
@classmethod
23-
def from_json(cls: Type['T'], obj: str) -> 'T':
24-
from google.protobuf import json_format
25-
from ...proto.docarray_pb2 import DocumentProto
40+
def from_json(
41+
cls: Type['T'], obj: str, protocol: str = 'jsonschema', **kwargs
42+
) -> 'T':
43+
"""Convert a JSON string into a Document.
44+
45+
:param obj: a valid JSON string
46+
:param protocol: `jsonschema` or `protobuf`
47+
:param kwargs: extra key-value args pass to pydantic and protobuf parser.
48+
:return: the parsed Document object
49+
"""
50+
if protocol == 'jsonschema':
51+
from ..pydantic_model import PydanticDocument
52+
53+
return cls.from_pydantic_model(PydanticDocument.parse_raw(obj, **kwargs))
54+
elif protocol == 'protobuf':
55+
from google.protobuf import json_format
56+
from ...proto.docarray_pb2 import DocumentProto
57+
58+
pb_msg = DocumentProto()
59+
json_format.Parse(obj, pb_msg, **kwargs)
60+
return cls.from_protobuf(pb_msg)
61+
else:
62+
raise ValueError(f'protocol=`{protocol}` is not supported')
2663

27-
pb_msg = DocumentProto()
28-
json_format.Parse(obj, pb_msg)
29-
return cls.from_protobuf(pb_msg)
64+
def to_dict(self, protocol: str = 'jsonschema', **kwargs) -> Dict[str, Any]:
65+
"""Convert itself into a Python dict object.
3066
31-
def to_dict(self, strict: bool = True) -> Dict[str, Any]:
32-
if strict:
67+
:param protocol: `jsonschema` or `protobuf`
68+
:param kwargs: extra key-value args pass to pydantic and protobuf dumper.
69+
:return: the dumped Document as a dict object
70+
"""
71+
if protocol == 'jsonschema':
72+
return self.to_pydantic_model().dict(**kwargs)
73+
elif protocol == 'protobuf':
3374
from google.protobuf.json_format import MessageToDict
3475

3576
return MessageToDict(
3677
self.to_protobuf(),
37-
preserving_proto_field_name=True,
78+
**kwargs,
3879
)
3980
else:
81+
warnings.warn(
82+
f'protocol=`{protocol}` is not supported, '
83+
f'the result dict is a Python dynamic typing dict without any promise on the schema.'
84+
)
4085
return dataclasses.asdict(self._data)
4186

4287
def to_bytes(
@@ -68,25 +113,33 @@ def from_bytes(
68113
"""
69114
bstr = decompress_bytes(data, algorithm=compress)
70115
if protocol == 'pickle':
71-
d = pickle.loads(bstr)
116+
return pickle.loads(bstr)
72117
elif protocol == 'protobuf':
73118
from ...proto.docarray_pb2 import DocumentProto
74119

75120
pb_msg = DocumentProto()
76121
pb_msg.ParseFromString(bstr)
77-
d = cls.from_protobuf(pb_msg)
122+
return cls.from_protobuf(pb_msg)
78123
else:
79124
raise ValueError(
80125
f'protocol={protocol} is not supported. Can be only `protobuf` or pickle protocols 0-5.'
81126
)
82-
return d
83127

84-
def to_json(self) -> str:
85-
from google.protobuf.json_format import MessageToJson
128+
def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str:
129+
"""Convert itself into a JSON string.
130+
131+
:param protocol: `jsonschema` or `protobuf`
132+
:param kwargs: extra key-value args pass to pydantic and protobuf dumper.
133+
:return: the dumped JSON string
134+
"""
135+
if protocol == 'jsonschema':
136+
return self.to_pydantic_model().json(**kwargs)
137+
elif protocol == 'protobuf':
138+
from google.protobuf.json_format import MessageToJson
86139

87-
return MessageToJson(
88-
self.to_protobuf(), preserving_proto_field_name=True, sort_keys=True
89-
)
140+
return MessageToJson(self.to_protobuf(), **kwargs)
141+
else:
142+
raise ValueError(f'protocol={protocol} is not supported.')
90143

91144
def to_base64(
92145
self, protocol: str = 'pickle', compress: Optional[str] = None

‎docarray/document/mixins/pydantic.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,11 @@ def from_pydantic_model(cls: Type['T'], model: 'BaseModel') -> 'T':
4747
from ... import Document
4848

4949
fields = {}
50+
_field_chunks, _field_matches = None, None
5051
if model.chunks:
51-
fields['chunks'] = [Document.from_pydantic_model(d) for d in model.chunks]
52+
_field_chunks = [Document.from_pydantic_model(d) for d in model.chunks]
5253
if model.matches:
53-
fields['matches'] = [Document.from_pydantic_model(d) for d in model.matches]
54+
_field_matches = [Document.from_pydantic_model(d) for d in model.matches]
5455

5556
for (field, value) in model.dict(
5657
exclude_none=True, exclude={'chunks', 'matches'}
@@ -66,4 +67,10 @@ def from_pydantic_model(cls: Type['T'], model: 'BaseModel') -> 'T':
6667
fields[f_name] = np.array(value)
6768
else:
6869
fields[f_name] = value
69-
return Document(**fields)
70+
71+
d = Document(**fields)
72+
if _field_chunks:
73+
d.chunks = _field_chunks
74+
if _field_matches:
75+
d.matches = _field_matches
76+
return d

‎docarray/document/pydantic_model.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@
1414

1515

1616
def _convert_ndarray_to_list(v: 'ArrayType'):
17-
return to_list(v)
17+
if v is not None:
18+
return to_list(v)
1819

1920

2021
class PydanticDocument(BaseModel):

0 commit comments

Comments
 (0)