attempt fix: tests

Author: aMahanna

arango/collection.py

@@ -169,14 +169,18 @@ def _prep_from_doc(
             else:
                 return doc_id, doc_id, {"If-Match": rev}
 
-    def _ensure_key_in_body(self, body: Json) -> Json:
-        """Return the document body with "_key" field populated.
+    def _ensure_key_in_body(self, body: Json, index: Optional[int] = None) -> Json:
+        """Return the document body with "_key" field populated. If
+            neither "_key" or "_id" exist, set "_key" value to **index**,
+            where **index** is the document's position in the sequence.
 
         :param body: Document body.
         :type body: dict
+        :param index: Document index value in the original list of documents.
+        :param index: int | None
         :return: Document body with "_key" field.
         :rtype: dict
-        :raise arango.exceptions.DocumentParseError: On missing ID and key.
+        :raise arango.exceptions.DocumentParseError: On missing _key, _id, & index.
         """
         if "_key" in body:
             return body
@@ -185,29 +189,24 @@ def _ensure_key_in_body(self, body: Json) -> Json:
             body = body.copy()
             body["_key"] = doc_id[len(self._id_prefix) :]
             return body
+        elif index:
+            body = body.copy()
+            body["_key"] = str(index)
+            return body
+
         raise DocumentParseError('field "_key" or "_id" required')
 
-    def _ensure_key_from_id(self, body: Json, index: Optional[int] = None) -> Json:
+    def _ensure_key_from_id(self, body: Json) -> Json:
         """Return the body with "_key" field if it has "_id" field.
-            If it has neither, set the "_key" value to i, where i
-            is the document's index position in the sequence.
-
         :param body: Document body.
         :type body: dict
-        :param index: Document index value in the original list of documents.
-        :param index: int | None
         :return: Document body with "_key" field if it has "_id" field.
         :rtype: dict
         """
         if "_id" in body and "_key" not in body:
             doc_id = self._validate_id(body["_id"])
             body = body.copy()
             body["_key"] = doc_id[len(self._id_prefix) :]
-
-        if "_id" not in body and "_key" not in body:
-            body = body.copy()
-            body["_key"] = str(index)
-
         return body
 
     @property
@@ -2002,7 +2001,7 @@ def import_bulk(
         :raise arango.exceptions.DocumentInsertError: If import fails.
         """
         documents = [
-            self._ensure_key_from_id(doc, i) for i, doc in enumerate(documents, 1)
+            self._ensure_key_in_body(doc, i) for i, doc in enumerate(documents, 1)
         ]
 
         params: Params = {"type": "array", "collection": self.name}

arango/utils.py

@@ -8,7 +8,7 @@
 
 import logging
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional, Union
+from typing import Any, Iterator, List, Optional, Sequence, Union
 
 from arango.exceptions import DocumentParseError
 from arango.typings import Json
@@ -85,8 +85,8 @@ def is_none_or_str(obj: Any) -> bool:
 
 
 def get_batches(
-    l: List[Any], batch_size: Optional[int] = None
-) -> Union[List[List[Any]], Iterator[List[Any]]]:
+    l: Sequence[Json], batch_size: Optional[int] = None
+) -> Union[List[Sequence[Json]], Iterator[Sequence[Json]]]:
     """Generator to split a list in batches
         of (maximum) **batch_size** elements each.
         If **batch_size** is invalid, return entire
@@ -100,7 +100,7 @@ def get_batches(
     if batch_size is None or batch_size <= 0 or batch_size >= len(l):
         return [l]
 
-    def generator() -> Iterator[List[Any]]:
+    def generator() -> Iterator[Sequence[Json]]:
         n = int(batch_size)  # type: ignore # (false positive)
         for i in range(0, len(l), n):
             yield l[i : i + n]