Merge pull request #618 from atihonruk/master

Implementation of the BZPOPMAX/BZPOPMIN commands and additional options of ZADD.

CHANGES/618.feature

@@ -0,0 +1,4 @@
+A few additions to the sorted set commands:
+
+- the blocking pop commands: `BZPOPMAX` and `BZPOPMIN`
+- the `CH` and `INCR` options of the `ZADD` command

CONTRIBUTORS.txt

@@ -31,6 +31,7 @@ Maxim Dodonchuk
 Michael Käufl
 Nickolai Novik
 Oleg Butuzov
+Oleksandr Tykhonruk
 Pau Freixes
 Paul Colomiets
 Samuel Colvin

aioredis/commands/sorted_set.py

@@ -18,7 +18,36 @@ class SortedSetCommandsMixin:
     ZSET_IF_NOT_EXIST = 'ZSET_IF_NOT_EXIST'  # NX
     ZSET_IF_EXIST = 'ZSET_IF_EXIST'  # XX
 
-    def zadd(self, key, score, member, *pairs, exist=None):
+    def bzpopmax(self, key, *keys, timeout=0, encoding=_NOTSET):
+        """Remove and get an element with the highest score in the sorted set,
+        or block until one is available.
+
+        :raises TypeError: if timeout is not int
+        :raises ValueError: if timeout is less than 0
+        """
+        if not isinstance(timeout, int):
+            raise TypeError("timeout argument must be int")
+        if timeout < 0:
+            raise ValueError("timeout must be greater equal 0")
+        args = keys + (timeout,)
+        return self.execute(b'BZPOPMAX', key, *args, encoding=encoding)
+
+    def bzpopmin(self, key, *keys, timeout=0, encoding=_NOTSET):
+        """Remove and get an element with the lowest score in the sorted set,
+        or block until one is available.
+
+        :raises TypeError: if timeout is not int
+        :raises ValueError: if timeout is less than 0
+        """
+        if not isinstance(timeout, int):
+            raise TypeError("timeout argument must be int")
+        if timeout < 0:
+            raise ValueError("timeout must be greater equal 0")
+        args = keys + (timeout,)
+        return self.execute(b'BZPOPMIN', key, *args, encoding=encoding)
+
+    def zadd(self, key, score, member, *pairs, exist=None, changed=False,
+             incr=False):
         """Add one or more members to a sorted set or update its score.
 
         :raises TypeError: score not int or float
@@ -39,6 +68,15 @@ def zadd(self, key, score, member, *pairs, exist=None):
         elif exist is self.ZSET_IF_NOT_EXIST:
             args.append(b'NX')
 
+        if changed:
+            args.append(b'CH')
+
+        if incr:
+            if pairs:
+                raise ValueError('only one score-element pair '
+                                 'can be specified in this mode')
+            args.append(b'INCR')
+
         args.extend([score, member])
         if pairs:
             args.extend(pairs)

tests/sorted_set_commands_test.py

@@ -5,6 +5,52 @@
 from _testutils import redis_version
 
 
+@redis_version(5, 0, 0, reason='BZPOPMAX is available since redis>=5.0.0')
+@pytest.mark.run_loop
+async def test_bzpopmax(redis):
+    key1 = b'key:zpopmax:1'
+    key2 = b'key:zpopmax:2'
+
+    pairs = [
+        (0, b'a'), (5, b'c'), (2, b'd'), (8, b'e'), (9, b'f'), (3, b'g')
+    ]
+    await redis.zadd(key1, *pairs[0])
+    await redis.zadd(key2, *itertools.chain.from_iterable(pairs))
+
+    res = await redis.bzpopmax(key1, timeout=0)
+    assert res == [key1, b'a', b'0']
+    res = await redis.bzpopmax(key1, key2, timeout=0)
+    assert res == [key2, b'f', b'9']
+
+    with pytest.raises(TypeError):
+        await redis.bzpopmax(key1, timeout=b'one')
+    with pytest.raises(ValueError):
+        await redis.bzpopmax(key2, timeout=-10)
+
+
+@redis_version(5, 0, 0, reason='BZPOPMIN is available since redis>=5.0.0')
+@pytest.mark.run_loop
+async def test_bzpopmin(redis):
+    key1 = b'key:zpopmin:1'
+    key2 = b'key:zpopmin:2'
+
+    pairs = [
+        (0, b'a'), (5, b'c'), (2, b'd'), (8, b'e'), (9, b'f'), (3, b'g')
+    ]
+    await redis.zadd(key1, *pairs[0])
+    await redis.zadd(key2, *itertools.chain.from_iterable(pairs))
+
+    res = await redis.bzpopmin(key1, timeout=0)
+    assert res == [key1, b'a', b'0']
+    res = await redis.bzpopmin(key1, key2, timeout=0)
+    assert res == [key2, b'a', b'0']
+
+    with pytest.raises(TypeError):
+        await redis.bzpopmin(key1, timeout=b'one')
+    with pytest.raises(ValueError):
+        await redis.bzpopmin(key2, timeout=-10)
+
+
 @pytest.mark.run_loop
 async def test_zadd(redis):
     key = b'key:zadd'
@@ -69,6 +115,15 @@ async def test_zadd_options(redis):
     res = await redis.zrange(key, 0, -1, withscores=False)
     assert res == [b'one', b'two']
 
+    res = await redis.zadd(key, 1, b'two', changed=True)
+    assert res == 1
+
+    res = await redis.zadd(key, 1, b'two', incr=True)
+    assert int(res) == 2
+
+    with pytest.raises(ValueError):
+        await redis.zadd(key, 1, b'one', 2, b'two', incr=True)
+
 
 @pytest.mark.run_loop
 async def test_zcard(redis):