add an optional argument to sort to make it return tuples

2 files changed, 61 insertions, 2 deletions

diff --git a/redis/client.py b/redis/client.py
index ebf9608..50c5905 100644
--- a/redis/client.py
+++ b/redis/client.py
@@ -674,7 +674,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, tuples=False):
         """
         Sort and return the list, set or sorted set at ``name``.
 
@@ -693,6 +693,11 @@ class StrictRedis(object):
 
         ``store`` allows for storing the result of the sort into
             the key ``store``
+
+        ``tuples`` if set to True and if ``get`` contains at least two
+        elements, will return a list of tuples matching the keys given
+        as arguments to ``get``.
+            
         """
         if (start is not None and num is None) or \
                 (num is not None and start is None):
@@ -725,7 +730,17 @@ class StrictRedis(object):
         if store is not None:
             pieces.append('STORE')
             pieces.append(store)
-        return self.execute_command('SORT', *pieces)
+
+        sorted_list = self.execute_command('SORT', *pieces)
+
+        if tuples:
+            if not get or isinstance(get, basestring) or len(get) < 2:
+                raise DataError("``get`` argument must be specified and "
+                                 "contain at least two keys")
+            n = len(get)
+            return zip(*[sorted_list[i::n] for i in range(n)])
+
+        return sorted_list
 
 
     #### SET COMMANDS ####
diff --git a/tests/server_commands.py b/tests/server_commands.py
index e1f90d3..9095fb1 100644
--- a/tests/server_commands.py
+++ b/tests/server_commands.py
@@ -1138,6 +1138,50 @@ class ServerCommandsTestCase(unittest.TestCase):
         self.assertEquals(self.client.sort('a', get=('user:*', '#')),
             ['u1', '1', 'u2', '2', 'u3', '3'])
 
+    def test_sort_get_tuples_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:*', '#'), tuples=True),
+            [('u1', '1'), ('u2', '2'), ('u3', '3')])
+
+    def test_sort_tuples_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:*', tuples=True)
+
+    def test_sort_tuples_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:*'], tuples=True)
+
+    def test_sort_tuples_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', tuples=True)
+
+    def test_sort_tuples_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:*', '#'), tuples=True),
+            [('u1', 'd1', '1'), ('u2', 'd2', '2'), ('u3', 'd3', '3')])
+
     def test_sort_desc(self):
         self.make_list('a', '231')
         self.assertEquals(self.client.sort('a', desc=True), ['3', '2', '1'])