diff --git a/CHANGES/618.feature b/CHANGES/618.feature new file mode 100644 index 000000000..f10db1430 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/CONTRIBUTORS.txt b/CONTRIBUTORS.txt index 53257133d..c668d8490 100644 --- a/CONTRIBUTORS.txt +++ b/CONTRIBUTORS.txt @@ -31,6 +31,7 @@ Maxim Dodonchuk Michael Käufl Nickolai Novik Oleg Butuzov +Oleksandr Tykhonruk Pau Freixes Paul Colomiets Samuel Colvin diff --git a/aioredis/commands/sorted_set.py b/aioredis/commands/sorted_set.py index fbc546b6d..1df2ab986 100644 --- a/aioredis/commands/sorted_set.py +++ b/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) diff --git a/tests/sorted_set_commands_test.py b/tests/sorted_set_commands_test.py index 84b47fc6d..b02119461 100644 --- a/tests/sorted_set_commands_test.py +++ b/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):