diff --git a/aioredis/commands/sorted_set.py b/aioredis/commands/sorted_set.py index 5ddf78d1d..fbc546b6d 100644 --- a/aioredis/commands/sorted_set.py +++ b/aioredis/commands/sorted_set.py @@ -424,6 +424,38 @@ def izscan(self, key, *, match=None, count=None): match=match, count=count)) + def zpopmin(self, key, count=None, *, encoding=_NOTSET): + """Removes and returns up to count members with the lowest scores + in the sorted set stored at key. + + :raises TypeError: if count is not int + """ + if count is not None and not isinstance(count, int): + raise TypeError("count argument must be int") + + args = [] + if count is not None: + args.extend([count]) + + fut = self.execute(b'ZPOPMIN', key, *args, encoding=encoding) + return fut + + def zpopmax(self, key, count=None, *, encoding=_NOTSET): + """Removes and returns up to count members with the highest scores + in the sorted set stored at key. + + :raises TypeError: if count is not int + """ + if count is not None and not isinstance(count, int): + raise TypeError("count argument must be int") + + args = [] + if count is not None: + args.extend([count]) + + fut = self.execute(b'ZPOPMAX', key, *args, encoding=encoding) + return fut + def _encode_min_max(flag, min, max): if flag is SortedSetCommandsMixin.ZSET_EXCLUDE_MIN: diff --git a/tests/sorted_set_commands_test.py b/tests/sorted_set_commands_test.py index cacda3971..84b47fc6d 100644 --- a/tests/sorted_set_commands_test.py +++ b/tests/sorted_set_commands_test.py @@ -1,4 +1,5 @@ import itertools + import pytest from _testutils import redis_version @@ -786,3 +787,37 @@ async def coro(cmd): with pytest.raises(TypeError): await redis.izscan(None) + + +@redis_version(5, 0, 0, reason='ZPOPMAX is available since redis>=5.0.0') +@pytest.mark.run_loop +async def test_zpopmax(redis): + key = b'key:zpopmax' + + pairs = [ + (0, b'a'), (5, b'c'), (2, b'd'), (8, b'e'), (9, b'f'), (3, b'g') + ] + await redis.zadd(key, *itertools.chain.from_iterable(pairs)) + + assert await redis.zpopmax(key) == [b'f', b'9'] + assert await redis.zpopmax(key, 3) == [b'e', b'8', b'c', b'5', b'g', b'3'] + + with pytest.raises(TypeError): + await redis.zpopmax(key, b'b') + + +@redis_version(5, 0, 0, reason='ZPOPMIN is available since redis>=5.0.0') +@pytest.mark.run_loop +async def test_zpopmin(redis): + key = b'key:zpopmin' + + pairs = [ + (0, b'a'), (5, b'c'), (2, b'd'), (8, b'e'), (9, b'f'), (3, b'g') + ] + await redis.zadd(key, *itertools.chain.from_iterable(pairs)) + + assert await redis.zpopmin(key) == [b'a', b'0'] + assert await redis.zpopmin(key, 3) == [b'd', b'2', b'g', b'3', b'c', b'5'] + + with pytest.raises(TypeError): + await redis.zpopmin(key, b'b')