diff --git a/redis/commands/core.py b/redis/commands/core.py index 383c635aea..266beb9a24 100644 --- a/redis/commands/core.py +++ b/redis/commands/core.py @@ -3259,6 +3259,31 @@ 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: 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 f4d7fa73e0..862f816de5 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2074,6 +2074,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=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):