Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit a96787e

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

13 files changed

+212-85Lines changed: 212 additions & 85 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎docarray/__init__.py‎

Copy file name to clipboard
+1-1Lines 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
Collapse file

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

Copy file name to clipboardExpand all lines: docarray/array/mixins/io/csv.py
+5-8Lines 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()})
Collapse file

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

Copy file name to clipboardExpand all lines: docarray/array/mixins/io/json.py
+28-16Lines 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
Collapse file

‎docarray/base.py‎

Copy file name to clipboardExpand all lines: docarray/base.py
+1-1Lines 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:
Collapse file

‎docarray/document/mixins/porting.py‎

Copy file name to clipboardExpand all lines: docarray/document/mixins/porting.py
+76-23Lines 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
Collapse file

‎docarray/document/mixins/pydantic.py‎

Copy file name to clipboardExpand all lines: docarray/document/mixins/pydantic.py
+10-3Lines 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
Collapse file

‎docarray/document/pydantic_model.py‎

Copy file name to clipboardExpand all lines: docarray/document/pydantic_model.py
+2-1Lines 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 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.