Add support for BZMPOP (redis#1851)

* add bzmpop

* add comment

* fix pr comment

* fix linters

* fix pr comments

Author: dvora-h

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
 
@@ -3237,6 +3238,38 @@ def bzpopmin(self, keys, timeout=0):
         keys.append(timeout)
         return self.execute_command("BZPOPMIN", *keys)
 
+    def bzmpop(
+        self,
+        timeout: float,
+        numkeys: int,
+        keys: List[str],
+        min: Optional[bool] = False,
+        max: Optional[bool] = False,
+        count: Optional[int] = 1,
+    ) -> Optional[list]:
+        """
+        Pop ``count`` values (default 1) off of the first non-empty sorted set
+        named in the ``keys`` list.
+
+        If none of the sorted sets in ``keys`` has a value to pop,
+        then block for ``timeout`` seconds, or until a member gets added
+        to one of the sorted sets.
+
+        If timeout is 0, then block indefinitely.
+
+        For more information check https://redis.io/commands/bzmpop
+        """
+        args = [timeout, numkeys, *keys]
+        if (min and max) or (not min and not max):
+            raise DataError("Either min or max, but not both must be set")
+        elif min:
+            args.append("MIN")
+        else:
+            args.append("MAX")
+        args.extend(["COUNT", count])
+
+        return self.execute_command("BZMPOP", *args)
+
     def _zrange(
         self,
         command,

tests/test_commands.py

@@ -2058,6 +2058,19 @@ 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_bzmpop(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.bzmpop(1, "2", ["b", "a"], min=True, count=2) == res
+        with pytest.raises(redis.DataError):
+            unstable_r.bzmpop(1, "2", ["b", "a"], count=2)
+        unstable_r.zadd("b", {"b1": 10, "ab": 9, "b3": 8})
+        res = [b"b", [[b"b1", b"10"]]]
+        assert unstable_r.bzmpop(0, "2", ["b", "a"], max=True) == res
+        assert unstable_r.bzmpop(1, "2", ["foo", "bar"], max=True) is None
+
     def test_zrange(self, r):
         r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
         assert r.zrange("a", 0, 1) == [b"a1", b"a2"]