Honour caller's 'limit' in iterview.

atomatt · atomatt · commit ba09f1c96f50 · 2013-04-25T11:50:03.000+02:00
diff --git a/couchdb/client.py b/couchdb/client.py
@@ -23,6 +23,7 @@
 >>> del server['python-tests']
 """
 
+import itertools
 import mimetypes
 import os
 from types import FunctionType
@@ -844,19 +845,30 @@ def iterview(self, name, batch, wrapper=None, **options):
         :param options: optional query string parameters
         :return: row generator
         """
+        # Check sane batch size.
         if batch <= 0:
             raise ValueError('batch must be 1 or more')
-        options['limit'] = batch + 1 # XXX todo: don't overwrite caller's limit
+        # Save caller's limit, it must be handled manually.
+        limit = options.get('limit')
+        if limit is not None and limit <= 0:
+            raise ValueError('limit must be 1 or more')
         startkey, startkey_docid = None, None # XXX todo: honour caller's startkey
         while True:
-            options.update(startkey=startkey, startkey_docid=startkey_docid)
-            # Get a batch of rows, with one extra for start of next batch.
+            loop_limit = min(limit or batch, batch)
+            # Get rows in batches, with one extra for start of next batch.
+            options.update(limit=loop_limit + 1, startkey=startkey,
+                           startkey_docid=startkey_docid)
             rows = list(self.view(name, wrapper, **options))
-            # Yield at most 'batch' rows.
-            for row in rows[:batch]:
+            # Yield rows from this batch.
+            for row in itertools.islice(rows, loop_limit):
                 yield row
-            if len(rows) <= batch:
+            # Decrement limit counter.
+            if limit is not None:
+                limit -= min(len(rows), batch)
+            # Check if there is nothing else to yield.
+            if len(rows) <= batch or (limit is not None and limit == 0):
                 break
+            # Save start keys for next loop.
             startkey, startkey_docid = rows[-1]['key'], rows[-1]['id']
 
     def show(self, name, docid=None, **options):
diff --git a/couchdb/tests/client.py b/couchdb/tests/client.py
@@ -730,7 +730,7 @@ def test_update_doc(self):
 
 class ViewIterationTestCase(testutil.TempDatabaseMixin, unittest.TestCase):
 
-    num_docs = 10
+    num_docs = 100
 
     def docfromnum(self, num):
         return {'_id': unicode(num), 'num': int(num / 2)}
@@ -763,6 +763,19 @@ def test_batchsizes(self):
         self.assertEqual(len(list(self.db.iterview('test/nums', self.num_docs))), self.num_docs)
         self.assertEqual(len(list(self.db.iterview('test/nums', self.num_docs + 1))), self.num_docs)
 
+    def test_limit(self):
+        # limit=0 doesn't make sense for iterview.
+        self.assertRaises(ValueError, self.db.iterview('test/nums', 10, limit=0).next)
+        # Test various limit sizes that are likely to cause trouble.
+        for limit in [1, int(self.num_docs / 4), self.num_docs - 1, self.num_docs,
+                      self.num_docs + 1]:
+            self.assertEqual([self.docfromrow(doc) for doc in self.db.iterview('test/nums', 10, limit=limit)],
+                             [self.docfromnum(x) for x in xrange(min(limit, self.num_docs))])
+        # Test limit same as batch size, in case of weird edge cases.
+        limit = int(self.num_docs / 4)
+        self.assertEqual([self.docfromrow(doc) for doc in self.db.iterview('test/nums', limit, limit=limit)],
+                         [self.docfromnum(x) for x in xrange(limit)])
+
 
 def suite():
     suite = unittest.TestSuite()
