diff options
| -rw-r--r-- | redis/commands/core.py | 25 | ||||
| -rw-r--r-- | tests/test_commands.py | 11 |
2 files changed, 36 insertions, 0 deletions
diff --git a/redis/commands/core.py b/redis/commands/core.py index 8c25979..80eb8e8 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3325,6 +3325,31 @@ class SortedSetCommands: keys.append(timeout) return self.execute_command("BZPOPMIN", *keys) + def zmpop( + self, + num_keys: int, + keys: List[str], + min: Optional[bool] = False, + max: Optional[bool] = False, + 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 + if (min and max) or (not min and not max): + raise DataError + elif min: + args.append("MIN") + else: + args.append("MAX") + if count != 1: + args.extend(["COUNT", count]) + + return self.execute_command("ZMPOP", *args) + def bzmpop( self, timeout: float, diff --git a/tests/test_commands.py b/tests/test_commands.py index c8f7ddc..a2d9335 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2108,6 +2108,17 @@ class TestRedisCommands: @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=True, count=2) == res + with pytest.raises(redis.DataError): + 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=True) == [b"b", [[b"b1", b"10"]]] + + @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"]]] |