Merge pull request #21 from novafloss/page-generator

Added a page generator

Author: brunobord

CHANGELOG.md

@@ -3,6 +3,7 @@
 # master (unreleased)
 
 - Deal with OneToOne Primary Keys (#19).
+- Expose the util function ``chunkator_page``.
 
 ## 1.1.0 (2017-03-25)
 

README.rst

@@ -44,8 +44,16 @@ You can also use ``values()``:
     you'll have several small queries. This will save your RAM instead, because
     you'll not load a huge queryset result before looping on it.
 
+If you want to manipulate the pages directly, you can use `chunkator_page`:
+
+.. code:: python
+
+    from chunkator import chunkator_page
+    queryset = LargeModel.objects.all().values('pk')
+    for page in chunkator_page(queryset, 200):
+        launch_some_task([item['pk'] for item in page])
 
 License
 =======
 
-MIT License.
+MIT License.

chunkator/__init__.py

@@ -12,13 +12,9 @@ class MissingPkFieldException(Exception):
     pass
 
 
-def chunkator(source_qs, chunk_size, query_log=None):
+def chunkator_page(source_qs, chunk_size, query_log=None):
     """
-    Yield over a queryset by chunks.
-
-    This method does not involve counting elements or measuring the iterable
-    length. We're saving at least a ``count()`` query on QuerySets, or a
-    CPU-and-RAM-consuming ``len(queryset)`` query.
+    Yield pages of a queryset.
     """
     pk = None
     # In django 1.9, _fields is always present and `None` if 'values()' is used
@@ -42,14 +38,35 @@ def chunkator(source_qs, chunk_size, query_log=None):
         page = queryset[:chunk_size]
         if query_log is not None:
             query_log.write('{page.query}\n'.format(page=page))
-        nb_items = 0
-        for item in page:
-            # source_qs._fields exists *and* is not none when using "values()"
-            if has_fields:
-                pk = item["pk"]
-            else:
-                pk = item.pk
-            nb_items += 1
-            yield item
+        page = list(page)
+        nb_items = len(page)
+
+        if nb_items == 0:
+            raise StopIteration
+
+        last_item = page[-1]
+        # source_qs._fields exists *and* is not none when using "values()"
+        if has_fields:
+            pk = last_item["pk"]
+        else:
+            pk = last_item.pk
+
+        yield page
+
         if nb_items < chunk_size:
             raise StopIteration
+
+
+def chunkator(source_qs, chunk_size, query_log=None):
+    """
+    Yield over a queryset by chunks.
+
+    This method does not involve counting elements or measuring the iterable
+    length. We're saving at least a ``count()`` query on QuerySets, or a
+    CPU-and-RAM-consuming ``len(queryset)`` query.
+    """
+    for page in chunkator_page(source_qs,
+                               chunk_size,
+                               query_log=query_log):
+        for item in page:
+            yield item

demo/demo_chunkator/tests.py

@@ -4,6 +4,7 @@
 
 from chunkator import MissingPkFieldException
 from chunkator import chunkator
+from chunkator import chunkator_page
 from demo_chunkator.models import Book
 from demo_chunkator.models import Cover
 from demo_chunkator.models import Profile
@@ -170,3 +171,18 @@ def test_query_log(self):
             # Should be 0 for the first query
             # Should occur once for other queries
             self.assertTrue(query.count('."uuid" >') <= 1, query)
+
+
+class ChunkatorPageTest(TestCase):
+    def setUp(self):
+        super(ChunkatorPageTest, self).setUp()
+        for nb in range(3):
+            Book.objects.create(
+                title="Title #{}".format(nb),
+                author="Author #{}".format(nb)
+            )
+
+    def test_chunkator_page(self):
+        qs = Book.objects.all().values('pk')
+        pages = list(chunkator_page(qs, 2))
+        self.assertEqual(pages, [[{'pk': 1}, {'pk': 2}], [{'pk': 3}]])