-
Notifications
You must be signed in to change notification settings - Fork 237
Expand file tree
/
Copy pathtypes.py
More file actions
299 lines (234 loc) · 8.9 KB
/
types.py
File metadata and controls
299 lines (234 loc) · 8.9 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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
import base64
import copy
import functools
import typing
from dataclasses import (
dataclass as _dataclass,
Field as _Field,
is_dataclass as _is_dataclass,
field as _field,
MISSING,
)
from enum import Enum
from typing import (
TypeVar,
Callable,
Optional,
overload,
Dict,
Type,
)
from docarray.dataclasses.getter import *
from docarray.dataclasses.setter import *
from docarray.dataclasses.enums import *
if TYPE_CHECKING: # pragma: no cover
import numpy as np
from docarray.typing import T
from docarray import Document
from docarray.typing import Image, Text, Audio, Video, Mesh, Tabular, Blob, JSON, URI
__all__ = ['field', 'dataclass', 'is_multimodal']
class AttributeTypeError(TypeError):
pass
class Field(_Field):
def __init__(
self,
*,
setter: Optional[Callable] = None,
getter: Optional[Callable] = None,
_source_field: Optional[_Field] = None,
**kwargs,
):
self.copy_from(_source_field if _source_field else _field(**kwargs))
self.setter = setter
self.getter = getter
def copy_from(self, f: '_Field'):
for s in f.__slots__:
setattr(self, s, getattr(f, s))
@overload
def field(
*,
_source_field: Optional[_Field] = None, # Privately used
setter: Optional[Callable] = None,
getter: Optional[Callable] = None,
default=MISSING,
default_factory=MISSING,
init=True,
repr=True,
hash=None,
compare=True,
metadata=None,
) -> Field:
...
def field(**kwargs) -> Field:
"""
Creates new multimodal type for a DocArray dataclass.
:meth:`field` is used to define the *get* and *set* behaviour of custom types when used in a DocArray dataclass.
.. code-block:: python
from docarray import Document, dataclass, field
from typing import TypeVar
MyImage = TypeVar('MyImage', bound=str)
def my_setter(value) -> 'Document':
return Document(uri=value).load_uri_to_blob()
def my_getter(doc: 'Document'):
return doc.uri
@dataclass
class MMDoc:
banner: MyImage = field(setter=my_setter, getter=my_getter, default='test-1.jpeg')
"""
return Field(**kwargs)
def _is_field(f) -> bool:
return isinstance(f, Field) and getattr(f, 'setter') and getattr(f, 'getter')
_TYPES_REGISTRY = {
Image: lambda x: field(setter=image_setter, getter=image_getter, _source_field=x),
Text: lambda x: field(setter=text_setter, getter=text_getter, _source_field=x),
URI: lambda x: field(setter=uri_setter, getter=uri_getter, _source_field=x),
Audio: lambda x: field(setter=audio_setter, getter=audio_getter, _source_field=x),
JSON: lambda x: field(setter=json_setter, getter=json_getter, _source_field=x),
Video: lambda x: field(setter=video_setter, getter=video_getter, _source_field=x),
Tabular: lambda x: field(
setter=tabular_setter, getter=tabular_getter, _source_field=x
),
Blob: lambda x: field(setter=blob_setter, getter=blob_getter, _source_field=x),
Mesh: lambda x: field(setter=mesh_setter, getter=mesh_getter, _source_field=x),
}
@overload
def dataclass(
cls: Optional['T'] = None,
*,
init: bool = True,
repr: bool = True,
eq: bool = True,
order: bool = False,
unsafe_hash: bool = False,
frozen: bool = False,
type_var_map: Optional[Dict[TypeVar, Callable[['_Field'], 'Field']]] = None,
) -> 'T':
...
def dataclass(
cls: Optional['T'] = None,
*,
type_var_map: Optional[Dict[TypeVar, Callable[['_Field'], 'Field']]] = None,
**kwargs,
) -> 'T':
"""Annotates a class as a DocArray dataclass type.
Example usage:
>>> from docarray.typing import Image, Text
>>> from docarray import dataclass
>>> @dataclass:
>>> class X:
>>> banner: Image = 'apple.png'
>>> description: Text = 'This is a big red apple.'
:param type_var_map: a mapping from TypeVar to a callable that gives Field.
.. highlight:: python
.. code-block:: python
_TYPES_REGISTRY = {
Image: lambda x: field(setter=image_setter, getter=image_getter, _source_field=x),
Text: lambda x: field(setter=text_setter, getter=text_getter, _source_field=x),
}
The default mapping will be overrided by this new mapping if they collide on the keys.
"""
if not type_var_map:
type_var_map = _TYPES_REGISTRY
else:
r = copy.deepcopy(_TYPES_REGISTRY)
r.update(type_var_map)
type_var_map = r
from docarray import Document
def deco(f):
"""
Set Decorator function.
:param f: function the decorator is used for
:return: wrapper
"""
@functools.wraps(f)
def wrapper(*args, **kwargs):
if not kwargs and len(args) == 2 and isinstance(args[1], Document):
return f(args[0], **_from_document(type(args[0]), args[1]))
else:
return f(*args, **kwargs)
return wrapper
def wrap(cls):
decorated_cls = _dataclass(cls, **kwargs)
# inject flag for recognizing this is a multimodal dataclass
setattr(decorated_cls, '__is_multimodal__', True)
# wrap init so `MMDoc(document)` is possible
if getattr(decorated_cls, '__init__'):
decorated_cls.__init__ = deco(decorated_cls.__init__)
for key, f in decorated_cls.__dataclass_fields__.items():
if _is_field(f):
continue
if f.type in type_var_map:
decorated_cls.__dataclass_fields__[key] = type_var_map[f.type](f)
elif isinstance(f.type, typing._GenericAlias) and f.type._name in [
'List',
'Iterable',
]:
sub_type = f.type.__args__[0]
if sub_type in type_var_map:
decorated_cls.__dataclass_fields__[key] = type_var_map[sub_type](f)
return decorated_cls
if cls is None:
return wrap
return wrap(cls)
def is_multimodal(obj) -> bool:
"""Returns True if obj is an instance of :meth:`dataclass`."""
from docarray import Document
if isinstance(obj, Document):
return obj.is_multimodal
else:
return _is_dataclass(obj) and hasattr(obj, '__is_multimodal__')
def _from_document(cls: Type['T'], doc: 'Document') -> 'T':
if not doc.is_multimodal:
raise ValueError(
f'{doc} is not a multimodal doc instantiated from a class wrapped by `docarray.dataclasses.tdataclass`.'
)
attributes = {}
for key, attribute_info in doc._metadata[
DocumentMetadata.MULTI_MODAL_SCHEMA
].items():
field = cls.__dataclass_fields__[key]
position = doc._metadata[DocumentMetadata.MULTI_MODAL_SCHEMA][key].get(
'position'
)
if (
attribute_info['type'] == 'bytes'
and attribute_info['attribute_type'] == AttributeType.PRIMITIVE
):
attributes[key] = base64.b64decode(doc.tags[key].encode())
elif attribute_info['attribute_type'] in [
AttributeType.PRIMITIVE,
AttributeType.ITERABLE_PRIMITIVE,
]:
attributes[key] = doc.tags[key]
elif attribute_info['attribute_type'] == AttributeType.DOCUMENT:
attribute_doc = doc.chunks[int(position)]
attribute = _get_doc_attribute(attribute_doc, field)
attributes[key] = attribute
elif attribute_info['attribute_type'] == AttributeType.ITERABLE_DOCUMENT:
attribute_list = []
for chunk_doc in doc.chunks[int(position)].chunks:
attribute_list.append(_get_doc_attribute(chunk_doc, field))
attributes[key] = attribute_list
elif attribute_info['attribute_type'] == AttributeType.NESTED:
nested_cls = field.type
attributes[key] = _get_doc_nested_attribute(
doc.chunks[int(position)], nested_cls
)
elif attribute_info['attribute_type'] == AttributeType.ITERABLE_NESTED:
nested_cls = cls.__dataclass_fields__[key].type.__args__[0]
attribute_list = []
for chunk_doc in doc.chunks[int(position)].chunks:
attribute_list.append(_get_doc_nested_attribute(chunk_doc, nested_cls))
attributes[key] = attribute_list
else:
raise AttributeError(f'Invalid attribute at `{key}`')
return attributes
def _get_doc_attribute(attribute_doc: 'Document', field):
if _is_field(field):
return field.getter(attribute_doc)
else:
raise ValueError('Invalid attribute type')
def _get_doc_nested_attribute(attribute_doc: 'Document', nested_cls: Type['T']) -> 'T':
if not is_multimodal(nested_cls):
raise ValueError(f'Nested attribute `{nested_cls.__name__}` is not a dataclass')
return nested_cls(**_from_document(nested_cls, attribute_doc))