Merge pull request #1037 from itamarhaber/zpop

Adds support for ZPOPMAX, ZPOPMIN and their respective blocking variants

1 files changed, 65 insertions, 1 deletions

diff --git a/redis/client.py b/redis/client.py
index a4d1cab..0cb6fca 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -448,9 +448,11 @@ class StrictRedis(object):
             lambda r: r and set(r) or set()
         ),
         string_keys_to_dict(
-            'ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE',
+            'ZPOPMAX ZPOPMIN ZRANGE ZRANGEBYSCORE ZREVRANGE ZREVRANGEBYSCORE',
             zset_score_pairs
         ),
+        string_keys_to_dict('BZPOPMIN BZPOPMAX', \
+                            lambda r: r and (r[0], r[1], float(r[2])) or None),
         string_keys_to_dict('ZRANK ZREVRANK', int_or_none),
         string_keys_to_dict('XREVRANGE XRANGE', parse_stream_list),
         string_keys_to_dict('XREAD XREADGROUP', parse_xread),
@@ -2115,6 +2117,68 @@ class StrictRedis(object):
         """
         return self.execute_command('ZLEXCOUNT', name, min, max)
 
+    def zpopmax(self, name, count=None):
+        """
+        Remove and return up to ``count`` members with the highest scores
+        from the sorted set ``name``.
+        """
+        args = (count is not None) and [count] or []
+        options = {
+            'withscores': True
+        }
+        return self.execute_command('ZPOPMAX', name, *args, **options)
+
+    def zpopmin(self, name, count=None):
+        """
+        Remove and return up to ``count`` members with the lowest scores
+        from the sorted set ``name``.
+        """
+        args = (count is not None) and [count] or []
+        options = {
+            'withscores': True
+        }
+        return self.execute_command('ZPOPMIN', name, *args, **options)
+
+    def bzpopmax(self, keys, timeout=0):
+        """
+        ZPOPMAX a value off of the first non-empty sorted set
+        named in the ``keys`` list.
+
+        If none of the sorted sets in ``keys`` has a value to ZPOPMAX,
+        then block for ``timeout`` seconds, or until a member gets added
+        to one of the sorted sets.
+
+        If timeout is 0, then block indefinitely.
+        """
+        if timeout is None:
+            timeout = 0
+        if isinstance(keys, basestring):
+            keys = [keys]
+        else:
+            keys = list(keys)
+        keys.append(timeout)
+        return self.execute_command('BZPOPMAX', *keys)
+
+    def bzpopmin(self, keys, timeout=0):
+        """
+        ZPOPMIN a value off of the first non-empty sorted set
+        named in the ``keys`` list.
+
+        If none of the sorted sets in ``keys`` has a value to ZPOPMIN,
+        then block for ``timeout`` seconds, or until a member gets added
+        to one of the sorted sets.
+
+        If timeout is 0, then block indefinitely.
+        """
+        if timeout is None:
+            timeout = 0
+        if isinstance(keys, basestring):
+            keys = [keys]
+        else:
+            keys = list(keys)
+        keys.append(timeout)
+        return self.execute_command('BZPOPMIN', *keys)
+
     def zrange(self, name, start, end, desc=False, withscores=False,
                score_cast_func=float):
         """