Merge pull request #1037 from itamarhaber/zpop

Adds support for ZPOPMAX, ZPOPMIN and their respective blocking variants

2 files changed, 115 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):
         """
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 22aafd3..579a85e 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -976,6 +976,56 @@ class TestRedisCommands(object):
         assert r.zrange('d', 0, -1, withscores=True) == \
             [(b('a3'), 20), (b('a1'), 23)]
 
+    @skip_if_server_version_lt('4.9.0')
+    def test_zpopmax(self, r):
+        r.delete('a')
+        r.zadd('a', a1=1, a2=2, a3=3)
+        assert r.zpopmax('a') == [(b('a3'), 3)]
+
+        # with count
+        assert r.zpopmax('a', count=2) == \
+            [(b('a2'), 2), (b('a1'), 1)]
+
+    @skip_if_server_version_lt('4.9.0')
+    def test_zpopmin(self, r):
+        r.delete('a')
+        r.zadd('a', a1=1, a2=2, a3=3)
+        assert r.zpopmin('a') == [(b('a1'), 1)]
+
+        # with count
+        assert r.zpopmin('a', count=2) == \
+            [(b('a2'), 2), (b('a3'), 3)]
+
+    @skip_if_server_version_lt('4.9.0')
+    def test_bzpopmax(self, r):
+        r.delete('a')
+        r.delete('b')
+        r.delete('c')
+        r.zadd('a', a1=1, a2=2)
+        r.zadd('b', b1=10, b2=20)
+        assert r.bzpopmax(['b', 'a'], timeout=1) == (b('b'), b('b2'), 20)
+        assert r.bzpopmax(['b', 'a'], timeout=1) == (b('b'), b('b1'), 10)
+        assert r.bzpopmax(['b', 'a'], timeout=1) == (b('a'), b('a2'), 2)
+        assert r.bzpopmax(['b', 'a'], timeout=1) == (b('a'), b('a1'), 1)
+        assert r.bzpopmax(['b', 'a'], timeout=1) is None
+        r.zadd('c', c1=100)
+        assert r.bzpopmax('c', timeout=1) == (b('c'), b('c1'), 100)
+
+    @skip_if_server_version_lt('4.9.0')
+    def test_bzpopmin(self, r):
+        r.delete('a')
+        r.delete('b')
+        r.delete('c')
+        r.zadd('a', a1=1, a2=2)
+        r.zadd('b', b1=10, b2=20)
+        assert r.bzpopmin(['b', 'a'], timeout=1) == (b('b'), b('b1'), 10)
+        assert r.bzpopmin(['b', 'a'], timeout=1) == (b('b'), b('b2'), 20)
+        assert r.bzpopmin(['b', 'a'], timeout=1) == (b('a'), b('a1'), 1)
+        assert r.bzpopmin(['b', 'a'], timeout=1) == (b('a'), b('a2'), 2)
+        assert r.bzpopmin(['b', 'a'], timeout=1) is None
+        r.zadd('c', c1=100)
+        assert r.bzpopmin('c', timeout=1) == (b('c'), b('c1'), 100)
+
     def test_zrange(self, r):
         r.zadd('a', a1=1, a2=2, a3=3)
         assert r.zrange('a', 0, 1) == [b('a1'), b('a2')]