fix: single doc set in docarray (#42)

davidbp · hanxiao · web-flow · commit 97725ed8c058 · 2022-01-13T12:27:59.000+01:00
Co-authored-by: Han Xiao &lt;han.xiao@jina.ai&gt;
diff --git a/docarray/array/document.py b/docarray/array/document.py
@@ -211,6 +211,7 @@ def __setitem__(
         index: 'DocumentArrayIndexType',
         value: Union['Document', Sequence['Document']],
     ):
+
         if isinstance(index, (int, np.generic)) and not isinstance(index, bool):
             index = int(index)
             self._data[index] = value
@@ -274,8 +275,11 @@ def __setitem__(
                         elif _a == 'embedding':
                             _docs.embeddings = _v
                         else:
-                            for _d, _vv in zip(_docs, _v):
-                                setattr(_d, _a, _vv)
+                            if len(_docs) == 1:
+                                setattr(_docs[0], _a, _v)
+                            else:
+                                for _d, _vv in zip(_docs, _v):
+                                    setattr(_d, _a, _vv)
             elif isinstance(index[0], bool):
                 if len(index) != len(self._data):
                     raise IndexError(
diff --git a/docarray/types.py b/docarray/types.py
@@ -48,8 +48,13 @@
     DocumentArrayMultipleIndexType = Union[
         slice, Sequence[int], Sequence[str], Sequence[bool], Ellipsis
     ]
-    DocumentArraySingleAttributeType = Tuple[slice, str]
-    DocumentArrayMultipleAttributeType = Tuple[slice, Sequence[str]]
+    DocumentArraySingleAttributeType = Tuple[
+        Union[DocumentArraySingletonIndexType, DocumentArrayMultipleIndexType], str
+    ]
+    DocumentArrayMultipleAttributeType = Tuple[
+        Union[DocumentArraySingletonIndexType, DocumentArrayMultipleIndexType],
+        Sequence[str],
+    ]
     DocumentArrayIndexType = Union[
         DocumentArraySingletonIndexType,
         DocumentArrayMultipleIndexType,
diff --git a/tests/unit/array/mixins/test_getset.py b/tests/unit/array/mixins/test_getset.py
@@ -94,6 +94,22 @@ def test_texts_getter_da(da):
     assert not da.texts
 
 
+@pytest.mark.parametrize('da', da_and_dam())
+def test_setter_by_sequences_in_selected_docs_da(da):
+
+    da[[0], 'text'] = 'jina'
+    assert ['jina'] == da[[0], 'text']
+
+    da[[0, 1], 'text'] = ['jina', 'jana']
+    assert ['jina', 'jana'] == da[[0, 1], 'text']
+
+    da[[0], 'id'] = '12'
+    assert ['12'] == da[[0], 'id']
+
+    da[[0, 1], 'id'] = ['12', '34']
+    assert ['12', '34'] == da[[0, 1], 'id']
+
+
 @pytest.mark.parametrize('da', da_and_dam())
 def test_texts_wrong_len(da):
     texts = ['hello']
