-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathtest_any_document.py
More file actions
91 lines (78 loc) · 2.53 KB
/
Copy pathtest_any_document.py
File metadata and controls
91 lines (78 loc) · 2.53 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
import numpy as np
import pytest
from typing import Dict, List
from docarray import DocList
from docarray.base_doc import AnyDoc, BaseDoc
from docarray.typing import NdArray
def test_any_doc():
class InnerDocument(BaseDoc):
text: str
tensor: NdArray
class CustomDoc(BaseDoc):
inner: InnerDocument
text: str
doc = CustomDoc(
text='bye', inner=InnerDocument(text='hello', tensor=np.zeros((3, 224, 224)))
)
any_doc = AnyDoc(**doc.__dict__)
assert any_doc.text == doc.text
assert any_doc.inner.text == doc.inner.text
assert (any_doc.inner.tensor == doc.inner.tensor).all()
@pytest.mark.parametrize('protocol', ['proto', 'json'])
def test_any_document_from_to(protocol):
class InnerDoc(BaseDoc):
text: str
t: Dict[str, str]
class DocTest(BaseDoc):
text: str
tags: Dict[str, int]
l: List[int]
d: InnerDoc
ld: DocList[InnerDoc]
inner_doc = InnerDoc(text='I am inner', t={'a': 'b'})
da = DocList[DocTest](
[
DocTest(
text='type1',
tags={'type': 1},
l=[1, 2],
d=inner_doc,
ld=DocList[InnerDoc]([inner_doc]),
),
DocTest(
text='type2',
tags={'type': 2},
l=[1, 2],
d=inner_doc,
ld=DocList[InnerDoc]([inner_doc]),
),
]
)
from docarray.base_doc import AnyDoc
if protocol == 'proto':
aux = DocList[AnyDoc].from_protobuf(da.to_protobuf())
else:
aux = DocList[AnyDoc].from_json(da.to_json())
assert len(aux) == 2
assert len(aux.id) == 2
for i, d in enumerate(aux):
assert d.tags['type'] == i + 1
assert d.text == f'type{i + 1}'
assert d.l == [1, 2]
if protocol == 'proto':
assert isinstance(d.d, AnyDoc)
assert d.d.text == 'I am inner' # inner Document is a Dict
assert d.d.t == {'a': 'b'}
else:
assert isinstance(d.d, dict)
assert d.d['text'] == 'I am inner' # inner Document is a Dict
assert d.d['t'] == {'a': 'b'}
assert len(d.ld) == 1
if protocol == 'proto':
assert isinstance(d.ld[0], AnyDoc)
assert d.ld[0].text == 'I am inner'
assert d.ld[0].t == {'a': 'b'}
else:
assert isinstance(d.ld[0], dict)
assert d.ld[0]['text'] == 'I am inner'
assert d.ld[0]['t'] == {'a': 'b'}