Merge pull request #83 from bigsassy/master

Ensure bulk_upsert accepts any iterable for rows

Author: Photonios

.gitignore

@@ -24,3 +24,6 @@ dist/
 # Ignore temporary tox environments
 .tox/
 .pytest_cache/
+
+# Ignore PyCharm / IntelliJ files
+.idea/

psqlextra/manager/manager.py

@@ -1,4 +1,5 @@
-from typing import Dict, List, Tuple, Union
+from typing import Dict, List, Tuple, Union, Iterable
+from itertools import chain
 
 import django
 
@@ -152,7 +153,7 @@ def bulk_insert(self, rows, return_model=False):
 
         Arguments:
             rows:
-                An array of dictionaries, where each dictionary
+                An iterable of dictionaries, where each dictionary
                 describes the fields to insert.
 
             return_model (default: False):
@@ -296,7 +297,7 @@ def upsert_and_get(
     def bulk_upsert(
         self,
         conflict_target: List,
-        rows: List[Dict],
+        rows: Iterable[Dict],
         index_predicate: str = None,
         return_model: bool = False,
     ):
@@ -322,21 +323,21 @@ def bulk_upsert(
             A list of either the dicts of the rows upserted, including the pk or
             the models of the rows upserted
         """
-
-        if not rows or len(rows) <= 0:
+        is_empty = lambda r: all([False for _ in r])
+        if not rows or is_empty(rows):
             return []
 
         self.on_conflict(
             conflict_target, ConflictAction.UPDATE, index_predicate
         )
         return self.bulk_insert(rows, return_model)
 
-    def _build_insert_compiler(self, rows: List[Dict]):
+    def _build_insert_compiler(self, rows: Iterable[Dict]):
         """Builds the SQL compiler for a insert query.
 
         Arguments:
             rows:
-                A list of dictionaries, where each entry
+                An iterable of dictionaries, where each entry
                 describes a record to insert.
 
         Returns:
@@ -349,8 +350,10 @@ def _build_insert_compiler(self, rows: List[Dict]):
         # we need to be certain that each row specifies the exact same
         # amount of fields/columns
         objs = []
-        field_count = len(rows[0])
-        for index, row in enumerate(rows):
+        rows_iter = iter(rows)
+        first_row = next(rows_iter)
+        field_count = len(first_row)
+        for index, row in enumerate(chain([first_row], rows_iter)):
             if field_count != len(row):
                 raise SuspiciousOperation(
                     (
@@ -366,7 +369,7 @@ def _build_insert_compiler(self, rows: List[Dict]):
         self._for_write = True
 
         # get the fields to be used during update/insert
-        insert_fields, update_fields = self._get_upsert_fields(rows[0])
+        insert_fields, update_fields = self._get_upsert_fields(first_row)
 
         # build a normal insert query
         query = PostgresInsertQuery(self.model)
@@ -597,7 +600,7 @@ def upsert_and_get(
     def bulk_upsert(
         self,
         conflict_target: List,
-        rows: List[Dict],
+        rows: Iterable[Dict],
         index_predicate: str = None,
         return_model: bool = False,
     ):

tests/test_upsert.py

@@ -146,3 +146,59 @@ def test_bulk_upsert_return_models():
     for index, obj in enumerate(objs, 1):
         assert isinstance(obj, model)
         assert obj.id == index
+
+
+def test_bulk_upsert_accepts_getitem_iterable():
+    """Tests whether an iterable only implementing
+    the __getitem__ method works correctly."""
+
+    class GetItemIterable:
+        def __init__(self, items):
+            self.items = items
+        def __getitem__(self, key):
+            return self.items[key]
+
+    model = get_fake_model(
+        {
+            "id": models.BigAutoField(primary_key=True),
+            "name": models.CharField(max_length=255, unique=True),
+        }
+    )
+
+    rows = GetItemIterable([dict(name="John Smith"), dict(name="Jane Doe")])
+
+    objs = model.objects.bulk_upsert(
+        conflict_target=["name"], rows=rows, return_model=True
+    )
+
+    for index, obj in enumerate(objs, 1):
+        assert isinstance(obj, model)
+        assert obj.id == index
+
+
+def test_bulk_upsert_accepts_iter_iterable():
+    """Tests whether an iterable only implementing
+    the __iter__ method works correctly."""
+
+    class IterIterable:
+        def __init__(self, items):
+            self.items = items
+        def __iter__(self):
+            return iter(self.items)
+
+    model = get_fake_model(
+        {
+            "id": models.BigAutoField(primary_key=True),
+            "name": models.CharField(max_length=255, unique=True),
+        }
+    )
+
+    rows = IterIterable([dict(name="John Smith"), dict(name="Jane Doe")])
+
+    objs = model.objects.bulk_upsert(
+        conflict_target=["name"], rows=rows, return_model=True
+    )
+
+    for index, obj in enumerate(objs, 1):
+        assert isinstance(obj, model)
+        assert obj.id == index