-
Notifications
You must be signed in to change notification settings - Fork 244
Expand file tree
/
Copy pathtest_base_getsetdel.py
More file actions
138 lines (101 loc) · 3.73 KB
/
test_base_getsetdel.py
File metadata and controls
138 lines (101 loc) · 3.73 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
from abc import ABC
from typing import Iterable, Sequence
import pytest
from docarray import DocumentArray, Document
from docarray.array.storage.base.getsetdel import BaseGetSetDelMixin
from docarray.array.storage.memory import BackendMixin, SequenceLikeMixin
class DummyGetSetDelMixin(BaseGetSetDelMixin):
"""Implement required and derived functions that power `getitem`, `setitem`, `delitem`"""
# essentials
def _del_doc_by_id(self, _id: str):
del self._data[self._id2offset[_id]]
self._id2offset.pop(_id)
def _del_doc_by_offset(self, offset: int):
self._id2offset.pop(self._data[offset].id)
del self._data[offset]
def _set_doc_by_id(self, _id: str, value: 'Document'):
old_idx = self._id2offset.pop(_id)
self._data[old_idx] = value
self._id2offset[value.id] = old_idx
def _get_doc_by_offset(self, offset: int) -> 'Document':
return self._data[offset]
def _get_doc_by_id(self, _id: str) -> 'Document':
return self._data[self._id2offset[_id]]
def _set_doc_by_offset(self, offset: int, value: 'Document'):
self._data[offset] = value
self._id2offset[value.id] = offset
def _get_docs_by_slice(self, _slice: slice) -> Iterable['Document']:
return self._data[_slice]
def _set_docs_by_slice(self, _slice: slice, value: Sequence['Document']):
self._data[_slice] = value
class StorageMixins(BackendMixin, DummyGetSetDelMixin, SequenceLikeMixin, ABC):
...
class DocumentArrayDummy(StorageMixins, DocumentArray):
def __new__(cls, *args, **kwargs):
return super().__new__(cls)
def _load_offset2ids(self):
pass
def _save_offset2ids(self):
pass
@pytest.fixture(scope='function')
def docs():
return DocumentArrayDummy([Document(id=str(j), text=str(j)) for j in range(100)])
def test_index_by_int_str(docs):
# getter
assert len(docs[[1]]) == 1
assert len(docs[1, 2]) == 2
assert len(docs[1, 2, 3]) == 3
assert len(docs[1:5]) == 4
assert len(docs[1:100:5]) == 20 # 1 to 100, sep with 5
# setter
with pytest.raises(TypeError, match='an iterable'):
docs[1:5] = Document(text='repl')
docs[1:5] = [Document(text=f'repl{j}') for j in range(4)]
for d in docs[1:5]:
assert d.text.startswith('repl')
assert len(docs) == 100
def test_getter_int_str(docs):
# getter
assert docs[99].text == '99'
assert docs[-1].text == '99'
assert docs[0].text == '0'
# string index
assert docs['0'].text == '0'
assert docs['99'].text == '99'
with pytest.raises(IndexError):
docs[100]
with pytest.raises(KeyError):
docs['adsad']
def test_set_content_none():
da = DocumentArray(
[
Document(mime_type='image'),
Document(mime_type='image'),
Document(mime_type='text'),
]
)
txt_da = da.find({'mime_type': {'$eq': 'image'}})
assert len(txt_da) == 2
txt_da.texts = ['hello', 'world']
assert txt_da.texts == ['hello', 'world']
assert da.texts == ['hello', 'world', '']
da.tensors = None
assert da.texts == ['hello', 'world', '']
def test_getter_from_docs():
da = DocumentArray.empty(4)
for i, d in enumerate(da):
d.id = f'root{i}'
chunks = DocumentArray.empty(2)
for j, c in enumerate(chunks):
c.id = f'chunk{j}_root{i}'
matches = DocumentArray.empty(2)
for j, m in enumerate(matches):
m.id = f'match{j}_root{i}'
d.chunks = chunks
d.matches = matches
assert len(da['@m']) == 8
assert len(da['@m[:3]']) == 8
assert len(da['@m[1:3]']) == 4
assert len(da[:, 'id']) == 4
assert len(da[...][:, 'id']) == 20
assert len(da[..., 'id']) == 20