Merge pull request #550 from leenr/master

Add support for `zpopmax` and `zpopmin` redis commands (server 5.0.0+)

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:

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')