-
Notifications
You must be signed in to change notification settings - Fork 238
Expand file tree
/
Copy pathporting.py
More file actions
88 lines (73 loc) · 2.71 KB
/
porting.py
File metadata and controls
88 lines (73 loc) · 2.71 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
import dataclasses
import pickle
from typing import Optional, TYPE_CHECKING, Type, Dict, Any
from ...helper import compress_bytes, decompress_bytes
if TYPE_CHECKING:
from ...types import T
class PortingMixin:
@classmethod
def from_dict(cls: Type['T'], obj: Dict) -> 'T':
from google.protobuf import json_format
from ...proto.docarray_pb2 import DocumentProto
pb_msg = DocumentProto()
json_format.ParseDict(obj, pb_msg)
return cls.from_protobuf(pb_msg)
@classmethod
def from_json(cls: Type['T'], obj: str) -> 'T':
from google.protobuf import json_format
from ...proto.docarray_pb2 import DocumentProto
pb_msg = DocumentProto()
json_format.Parse(obj, pb_msg)
return cls.from_protobuf(pb_msg)
def to_dict(self, strict: bool = True) -> Dict[str, Any]:
if strict:
from google.protobuf.json_format import MessageToDict
return MessageToDict(
self.to_protobuf(),
preserving_proto_field_name=True,
)
else:
return dataclasses.asdict(self._data)
def to_bytes(
self, protocol: str = 'pickle', compress: Optional[str] = None
) -> bytes:
if protocol == 'pickle':
bstr = pickle.dumps(self)
elif protocol == 'protobuf':
bstr = self.to_protobuf().SerializePartialToString()
else:
raise ValueError(
f'protocol={protocol} is not supported. Can be only `protobuf` or pickle protocols 0-5.'
)
return compress_bytes(bstr, algorithm=compress)
@classmethod
def from_bytes(
cls: Type['T'],
data: bytes,
protocol: str = 'pickle',
compress: Optional[str] = None,
) -> 'T':
"""Build Document object from binary bytes
:param data: binary bytes
:param protocol: protocol to use
:param compress: compress method to use
:return: a Document object
"""
bstr = decompress_bytes(data, algorithm=compress)
if protocol == 'pickle':
d = pickle.loads(bstr)
elif protocol == 'protobuf':
from ...proto.docarray_pb2 import DocumentProto
pb_msg = DocumentProto()
pb_msg.ParseFromString(bstr)
d = cls.from_protobuf(pb_msg)
else:
raise ValueError(
f'protocol={protocol} is not supported. Can be only `protobuf` or pickle protocols 0-5.'
)
return d
def to_json(self) -> str:
from google.protobuf.json_format import MessageToJson
return MessageToJson(
self.to_protobuf(), preserving_proto_field_name=True, sort_keys=True
)