From 3ae42ad9eb2284696bf9f1626ffcd0d7b383ce8e Mon Sep 17 00:00:00 2001
From: Janne Vanhala <janne.vanhala@gmail.com>
Date: Mon, 25 Aug 2014 17:17:21 +0300
Subject: [PATCH] Add `count()` method to `QueryChain`

---
 sqlalchemy_utils/query_chain.py | 31 ++++++++++++++++++++++++++++++-
 tests/test_query_chain.py       |  3 +++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/sqlalchemy_utils/query_chain.py b/sqlalchemy_utils/query_chain.py
index cd2d5a1..1596b43 100644
--- a/sqlalchemy_utils/query_chain.py
+++ b/sqlalchemy_utils/query_chain.py
@@ -5,7 +5,7 @@ QueryChain is a wrapper for sequence of queries.
 Features:
 
     * Easy iteration for sequence of queries
-    * Limit and offset which are applied to all queries in the chain
+    * Limit, offset and count which are applied to all queries in the chain
     * Smart __getitem__ support
 
 
@@ -81,6 +81,29 @@ Chain slicing
     )
 
     chain[3:6]   # New QueryChain with offset=3 and limit=6
+
+
+Count
+^^^^^
+
+Let's assume that there are five blog posts, five articles and five news
+items in the database, and you have the following query chain::
+
+    chain = QueryChain(
+        [
+            session.query(BlogPost),
+            session.query(Article),
+            session.query(NewsItem)
+        ]
+    )
+
+You can then get the total number rows returned by the query chain
+with :meth:`~QueryChain.count`::
+
+    >>> chain.count()
+    15
+
+
 """
 from copy import copy
 
@@ -129,6 +152,12 @@ class QueryChain(object):
     def offset(self, value):
         return self[value:]
 
+    def count(self):
+        """
+        Return the total number of rows this QueryChain's queries would return.
+        """
+        return sum(q.count() for q in self.queries)
+
     def __getitem__(self, key):
         if isinstance(key, slice):
             return self.__class__(
diff --git a/tests/test_query_chain.py b/tests/test_query_chain.py
index c6c1613..7796ff7 100644
--- a/tests/test_query_chain.py
+++ b/tests/test_query_chain.py
@@ -88,3 +88,6 @@ class TestQueryChain(TestCase):
     def test_getitem_with_single_key(self):
         article = self.chain[2]
         assert article == self.articles[0]
+
+    def test_count(self):
+        assert self.chain.count() == 9