Merge branch 'mapleoin-master'

3 files changed, 79 insertions, 2 deletions

diff --git a/CHANGES b/CHANGES
index 3134d2e..c845202 100644
--- a/CHANGES
+++ b/CHANGES
@@ -28,6 +28,10 @@
     * Added PSETEX. Thanks YAMAMOTO Takashi.
     * Added a BlockingConnectionPool to limit the number of connections that
       can be created. Thanks James Arthur.
+    * SORT now accepts a `groups` option that if specified, will return
+      tuples of n-length, where n is the number of keys specified in the GET
+      argument. This allows for convenient row-based iteration. Thanks
+      Ionuț Arțăriși.
 * 2.7.2
     * Parse errors are now *always* raised on multi/exec pipelines, regardless
       of the `raise_on_error` flag. See
diff --git a/redis/client.py b/redis/client.py
index d9506e6..ea8b5cf 100644
--- a/redis/client.py
+++ b/redis/client.py
@@ -127,6 +127,16 @@ def zset_score_pairs(response, **options):
     return list(izip(it, imap(score_cast_func, it)))
 
 
+def sort_return_tuples(response, **options):
+    """
+    If ``groups`` is specified, return the response as a list of
+    n-element tuples with n being the value found in options['groups']
+    """
+    if not response or not options['groups']:
+        return response
+    n = options['groups']
+    return list(izip(*[response[i::n] for i in range(n)]))
+
 def int_or_none(response):
     if response is None:
         return None
@@ -198,6 +208,7 @@ class StrictRedis(object):
             'LPUSH RPUSH',
             lambda r: isinstance(r, long) and r or nativestr(r) == 'OK'
         ),
+        string_keys_to_dict('SORT', sort_return_tuples),
         string_keys_to_dict('ZSCORE ZINCRBY', float_or_none),
         string_keys_to_dict(
             'FLUSHALL FLUSHDB LSET LTRIM MSET RENAME '
@@ -912,7 +923,7 @@ class StrictRedis(object):
         return self.execute_command('RPUSHX', name, value)
 
     def sort(self, name, start=None, num=None, by=None, get=None,
-             desc=False, alpha=False, store=None):
+             desc=False, alpha=False, store=None, groups=False):
         """
         Sort and return the list, set or sorted set at ``name``.
 
@@ -931,6 +942,11 @@ class StrictRedis(object):
 
         ``store`` allows for storing the result of the sort into
             the key ``store``
+
+        ``groups`` if set to True and if ``get`` contains at least two
+            elements, sort will return a list of tuples, each containing the
+            values fetched from the arguments to ``get``.
+
         """
         if (start is not None and num is None) or \
                 (num is not None and start is None):
@@ -963,7 +979,15 @@ class StrictRedis(object):
         if store is not None:
             pieces.append('STORE')
             pieces.append(store)
-        return self.execute_command('SORT', *pieces)
+
+        if groups:
+            if not get or isinstance(get, basestring) or len(get) < 2:
+                raise DataError('when using "groups" the "get" argument '
+                                'must be specified and contain at least '
+                                'two keys')
+
+        options = {'groups': len(get) if groups else None}
+        return self.execute_command('SORT', *pieces, **options)
 
     #### SET COMMANDS ####
     def sadd(self, name, *values):
diff --git a/tests/server_commands.py b/tests/server_commands.py
index 17adcf8..0225242 100644
--- a/tests/server_commands.py
+++ b/tests/server_commands.py
@@ -1513,6 +1513,55 @@ class ServerCommandsTestCase(unittest.TestCase):
             self.client.sort('a', get=('user:*', '#')),
             [b('u1'), b('1'), b('u2'), b('2'), b('u3'), b('3')])
 
+    def test_sort_get_groups_two(self):
+        self.client['user:1'] = 'u1'
+        self.client['user:2'] = 'u2'
+        self.client['user:3'] = 'u3'
+        self.make_list('a', '231')
+        self.assertEquals(
+            self.client.sort('a', get=('user:*', '#'), groups=True),
+            [(b('u1'), b('1')), (b('u2'), b('2')), (b('u3'), b('3'))])
+
+    def test_sort_groups_string_get(self):
+        self.client['user:1'] = 'u1'
+        self.client['user:2'] = 'u2'
+        self.client['user:3'] = 'u3'
+        self.make_list('a', '231')
+        self.assertRaises(redis.DataError, self.client.sort, 'a',
+                          get='user:*', groups=True)
+
+    def test_sort_groups_just_one_get(self):
+        self.client['user:1'] = 'u1'
+        self.client['user:2'] = 'u2'
+        self.client['user:3'] = 'u3'
+        self.make_list('a', '231')
+        self.assertRaises(redis.DataError, self.client.sort, 'a',
+                          get=['user:*'], groups=True)
+
+    def test_sort_groups_no_get(self):
+        self.client['user:1'] = 'u1'
+        self.client['user:2'] = 'u2'
+        self.client['user:3'] = 'u3'
+        self.make_list('a', '231')
+        self.assertRaises(redis.DataError, self.client.sort, 'a', groups=True)
+
+    def test_sort_groups_three_gets(self):
+        self.client['user:1'] = 'u1'
+        self.client['user:2'] = 'u2'
+        self.client['user:3'] = 'u3'
+        self.client['door:1'] = 'd1'
+        self.client['door:2'] = 'd2'
+        self.client['door:3'] = 'd3'
+        self.make_list('a', '231')
+        self.assertEquals(
+            self.client.sort('a', get=('user:*', 'door:*', '#'), groups=True),
+            [
+                (b('u1'), b('d1'), b('1')),
+                (b('u2'), b('d2'), b('2')),
+                (b('u3'), b('d3'), b('3'))
+            ]
+        )
+
     def test_sort_desc(self):
         self.make_list('a', '231')
         self.assertEquals(