Add support for ZMPOP

redis/commands/core.py

@@ -2,6 +2,7 @@
 import hashlib
 import time
 import warnings
+from typing import List, Optional
 
 from redis.exceptions import ConnectionError, DataError, NoScriptError, RedisError
 
@@ -3255,6 +3256,25 @@ def bzpopmin(self, keys, timeout=0):
         keys.append(timeout)
         return self.execute_command("BZPOPMIN", *keys)
 
+    def zmpop(
+        self,
+        num_keys: int,
+        keys: List[str],
+        min_max: str,
+        count: Optional[int] = 1,
+    ) -> list:
+        """
+        Pop ``count`` values (default 1) off of the first non-empty sorted set
+        named in the ``keys`` list.
+
+        For more information check https://redis.io/commands/zmpop
+        """
+        args = [num_keys] + keys + [min_max]
+        if count != 1:
+            args.extend(["COUNT", count])
+
+        return self.execute_command("ZMPOP", *args)
+
     def _zrange(
         self,
         command,

tests/test_commands.py

@@ -2058,6 +2058,17 @@ def test_bzpopmin(self, r):
         r.zadd("c", {"c1": 100})
         assert r.bzpopmin("c", timeout=1) == (b"c", b"c1", 100)
 
+    @pytest.mark.onlynoncluster
+    # @skip_if_server_version_lt("7.0.0") turn on after redis 7 release
+    def test_zmpop(self, unstable_r):
+        unstable_r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
+        res = [b"a", [[b"a1", b"1"], [b"a2", b"2"]]]
+        assert unstable_r.zmpop("2", ["b", "a"], "MIN", count=2) == res
+        with pytest.raises(TypeError):
+            unstable_r.zmpop("2", ["b", "a"], count=2)
+        unstable_r.zadd("b", {"b1": 10, "ab": 9, "b3": 8})
+        assert unstable_r.zmpop("2", ["b", "a"], "MAX") == [b"b", [[b"b1", b"10"]]]
+
     def test_zrange(self, r):
         r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
         assert r.zrange("a", 0, 1) == [b"a1", b"a2"]