diff --git a/redis/commands/core.py b/redis/commands/core.py index 4f0accd957..6b3bc4e76b 100644 --- a/redis/commands/core.py +++ b/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,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, diff --git a/tests/test_commands.py b/tests/test_commands.py index b28b63ea6e..937172a195 100644 --- a/tests/test_commands.py +++ b/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"]