Add support for ZMPOP (#1923)

Author: dvora-h

redis/commands/core.py

@@ -3325,6 +3325,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,

tests/test_commands.py

@@ -2106,6 +2106,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):