feat: refactor chunkset append method

jina-ai · Nov 29, 2020 · b6c24f5 · b6c24f5
1 parent 06284b1
commit b6c24f5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/jina/types/sets/chunk_set.py b/jina/types/sets/chunk_set.py
@@ -21,15 +21,16 @@ def append(self, document: Optional['Document'] = None, **kwargs) -> 'Document':
             make sure the added chunk is legit.
         """
         from ..document import Document
-        chunk = Document()
-        if document:
-            chunk.CopyFrom(document)
-
-        chunk.set_attrs(parent_id=self._ref_doc.id,
-                        granularity=self._ref_doc.granularity + 1,
-                        **kwargs)
-        if not chunk.mime_type:
-            chunk.mime_type = self._ref_doc.mime_type
-
-        self._docs_proto.append(chunk.as_pb_object)
-        return chunk
+        with Document() as chunk:
+            if document:
+                chunk.CopyFrom(document)
+
+            chunk.set_attrs(parent_id=self._ref_doc.id,
+                            granularity=self._ref_doc.granularity + 1,
+                            **kwargs)
+
+            if not chunk.mime_type:
+                chunk.mime_type = self._ref_doc.mime_type
+
+            self._docs_proto.append(chunk.as_pb_object)
+            return chunk
diff --git a/tests/unit/types/test_chunkset.py b/tests/unit/types/test_chunkset.py
@@ -36,6 +36,5 @@ def test_append_from_documents(chunkset, document_factory, reference_doc):
     assert len(chunkset) == 4
     assert chunkset[-1].text == 'test 4'
     assert rv.text == chunk.text
-    assert rv.id == chunk.id
     assert rv.parent_id == reference_doc.id
     assert rv.granularity == reference_doc.granularity + 1