From 64a31b4478ff4078f97f7acc55fee9a67fe7dedb Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Wed, 7 May 2014 09:39:19 +0300
Subject: [PATCH] Add chained method call API for QueryChain

---
 sqlalchemy_utils/query_chain.py | 32 +++++++++++++++++++++++---------
 tests/test_query_chain.py       | 21 ++++++++++-----------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/sqlalchemy_utils/query_chain.py b/sqlalchemy_utils/query_chain.py
index 6b879bd..cd2d5a1 100644
--- a/sqlalchemy_utils/query_chain.py
+++ b/sqlalchemy_utils/query_chain.py
@@ -55,10 +55,18 @@ database.
     list(chain)  # all blog posts but not articles and news items
 
 
-    chain.offset = 4
+    chain = chain.offset(4)
     list(chain)  # last blog post, and first four articles
 
 
+Just like with original query object the limit and offset can be chained to
+return a new QueryChain.
+
+::
+
+    chain = chain.limit(5).offset(7)
+
+
 Chain slicing
 ^^^^^^^^^^^^^
 
@@ -91,18 +99,18 @@ class QueryChain(object):
     """
     def __init__(self, queries, limit=None, offset=None):
         self.queries = queries
-        self.limit = limit
-        self.offset = offset
+        self._limit = limit
+        self._offset = offset
 
     def __iter__(self):
         consumed = 0
         skipped = 0
         for query in self.queries:
             query_copy = copy(query)
-            if self.limit:
-                query = query.limit(self.limit - consumed)
-            if self.offset:
-                query = query.offset(self.offset - skipped)
+            if self._limit:
+                query = query.limit(self._limit - consumed)
+            if self._offset:
+                query = query.offset(self._offset - skipped)
 
             obj_count = 0
             for obj in query:
@@ -115,12 +123,18 @@ class QueryChain(object):
             else:
                 skipped += obj_count
 
+    def limit(self, value):
+        return self[:value]
+
+    def offset(self, value):
+        return self[value:]
+
     def __getitem__(self, key):
         if isinstance(key, slice):
             return self.__class__(
                 queries=self.queries,
-                limit=key.stop,
-                offset=key.start
+                limit=key.stop if key.stop is not None else self._limit,
+                offset=key.start if key.start is not None else self._offset
             )
         else:
             for obj in self[key:1]:
diff --git a/tests/test_query_chain.py b/tests/test_query_chain.py
index ebae8ae..c6c1613 100644
--- a/tests/test_query_chain.py
+++ b/tests/test_query_chain.py
@@ -57,25 +57,24 @@ class TestQueryChain(TestCase):
         assert len(list(self.chain)) == 9
 
     def test_iter_with_limit(self):
-        self.chain.limit = 4
-        objects = list(self.chain)
+        chain = self.chain.limit(4)
+        objects = list(chain)
         assert self.users == objects[0:2]
         assert self.articles[0:2] == objects[2:]
 
     def test_iter_with_offset(self):
-        self.chain.offset = 3
-        objects = list(self.chain)
+        chain = self.chain.offset(3)
+        objects = list(chain)
         assert self.articles[1:] + self.posts == objects
 
     def test_iter_with_limit_and_offset(self):
-        self.chain.offset = 3
-        self.chain.limit = 4
-        objects = list(self.chain)
+        chain = self.chain.offset(3).limit(4)
+        objects = list(chain)
         assert self.articles[1:] + self.posts[0:1] == objects
 
     def test_iter_with_offset_spanning_multiple_queries(self):
-        self.chain.offset = 7
-        objects = list(self.chain)
+        chain = self.chain.offset(7)
+        objects = list(chain)
         assert self.posts[1:] == objects
 
     def test_repr(self):
@@ -83,8 +82,8 @@ class TestQueryChain(TestCase):
 
     def test_getitem_with_slice(self):
         chain = self.chain[1:]
-        assert chain.offset == 1
-        assert chain.limit is None
+        assert chain._offset == 1
+        assert chain._limit is None
 
     def test_getitem_with_single_key(self):
         article = self.chain[2]