Python: add ZDIFF command (valkey-io#1401)

* Python: add ZDIFF command (#255)

* Add PR link in CHANGELOG

* PR suggestions

CHANGELOG.md

@@ -24,7 +24,7 @@
 * Python: Added BLPOP and BRPOP commands ([#1369](https://github.com/aws/glide-for-redis/pull/1369))
 * Python: Added ZRANGESTORE command ([#1377](https://github.com/aws/glide-for-redis/pull/1377))
 * Python: Added ZDIFFSTORE command ([#1378](https://github.com/aws/glide-for-redis/pull/1378))
-
+* Python: Added ZDIFF command ([#1401](https://github.com/aws/glide-for-redis/pull/1401))
 
 #### Fixes
 * Python: Fix typing error "‘type’ object is not subscriptable" ([#1203](https://github.com/aws/glide-for-redis/pull/1203))

python/python/glide/async_commands/core.py

@@ -2712,6 +2712,65 @@ async def zmscore(
             await self._execute_command(RequestType.ZMScore, [key] + members),
         )
 
+    async def zdiff(self, keys: List[str]) -> List[str]:
+        """
+        Returns the difference between the first sorted set and all the successive sorted sets.
+        To get the elements with their scores, see `zdiff_withscores`.
+
+        When in Cluster mode, all keys must map to the same hash slot.
+
+        See https://valkey.io/commands/zdiff for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Returns:
+            List[str]: A list of elements representing the difference between the sorted sets.
+                If the first key does not exist, it is treated as an empty sorted set, and the command returns an
+                empty list.
+
+        Examples:
+            >>> await client.zadd("sorted_set1", {"element1":1.0, "element2": 2.0, "element3": 3.0})
+            >>> await client.zadd("sorted_set2", {"element2": 2.0})
+            >>> await client.zadd("sorted_set3", {"element3": 3.0})
+            >>> await client.zdiff("sorted_set1", "sorted_set2", "sorted_set3")
+                ["element1"]  # Indicates that "element1" is in "sorted_set1" but not "sorted_set2" or "sorted_set3".
+        """
+        return cast(
+            List[str],
+            await self._execute_command(RequestType.ZDiff, [str(len(keys))] + keys),
+        )
+
+    async def zdiff_withscores(self, keys: List[str]) -> Mapping[str, float]:
+        """
+        Returns the difference between the first sorted set and all the successive sorted sets, with the associated scores.
+        When in Cluster mode, all keys must map to the same hash slot.
+
+        See https://valkey.io/commands/zdiff for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Returns:
+            Mapping[str, float]: A dictionary of elements and their scores representing the difference between the sorted
+                sets.
+                If the first `key` does not exist, it is treated as an empty sorted set, and the command returns an
+                empty list.
+
+        Examples:
+            >>> await client.zadd("sorted_set1", {"element1":1.0, "element2": 2.0, "element3": 3.0})
+            >>> await client.zadd("sorted_set2", {"element2": 2.0})
+            >>> await client.zadd("sorted_set3", {"element3": 3.0})
+            >>> await client.zdiff_withscores("sorted_set1", "sorted_set2", "sorted_set3")
+                {"element1": 1.0}  # Indicates that "element1" is in "sorted_set1" but not "sorted_set2" or "sorted_set3".
+        """
+        return cast(
+            Mapping[str, float],
+            await self._execute_command(
+                RequestType.ZDiff, [str(len(keys))] + keys + ["WITHSCORES"]
+            ),
+        )
+
     async def zdiffstore(self, destination: str, keys: List[str]) -> int:
         """
         Calculates the difference between the first sorted set and all the successive sorted sets at `keys` and stores

python/python/glide/async_commands/transaction.py

@@ -1945,6 +1945,41 @@ def zmscore(self: TTransaction, key: str, members: List[str]) -> TTransaction:
         """
         return self.append_command(RequestType.ZMScore, [key] + members)
 
+    def zdiff(self: TTransaction, keys: List[str]) -> TTransaction:
+        """
+        Returns the difference between the first sorted set and all the successive sorted sets.
+        To get the elements with their scores, see `zdiff_withscores`.
+
+        See https://valkey.io/commands/zdiff for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Command response:
+            List[str]: A list of elements representing the difference between the sorted sets.
+                If the first key does not exist, it is treated as an empty sorted set, and the command returns an
+                empty list.
+        """
+        return self.append_command(RequestType.ZDiff, [str(len(keys))] + keys)
+
+    def zdiff_withscores(self: TTransaction, keys: List[str]) -> TTransaction:
+        """
+        Returns the difference between the first sorted set and all the successive sorted sets, with the associated scores.
+
+        See https://valkey.io/commands/zdiff for more details.
+
+        Args:
+            keys (List[str]): The keys of the sorted sets.
+
+        Command response:
+            Mapping[str, float]: A dictionary of elements and their scores representing the difference between the sorted sets.
+                If the first `key` does not exist, it is treated as an empty sorted set, and the command returns an
+                empty list.
+        """
+        return self.append_command(
+            RequestType.ZDiff, [str(len(keys))] + keys + ["WITHSCORES"]
+        )
+
     def zdiffstore(
         self: TTransaction, destination: str, keys: List[str]
     ) -> TTransaction:

python/python/tests/test_async_client.py

@@ -2332,6 +2332,70 @@ async def test_zrank(self, redis_client: TRedisClient):
         with pytest.raises(RequestError):
             await redis_client.zrank(key, "one")
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_zdiff(self, redis_client: TRedisClient):
+        key1 = f"{{testKey}}:1-{get_random_string(10)}"
+        key2 = f"{{testKey}}:2-{get_random_string(10)}"
+        key3 = f"{{testKey}}:3-{get_random_string(10)}"
+        string_key = f"{{testKey}}:4-{get_random_string(10)}"
+        non_existing_key = f"{{testKey}}:5-{get_random_string(10)}"
+
+        member_scores1 = {"one": 1.0, "two": 2.0, "three": 3.0}
+        member_scores2 = {"two": 2.0}
+        member_scores3 = {"one": 1.0, "two": 2.0, "three": 3.0, "four": 4.0}
+
+        assert await redis_client.zadd(key1, member_scores1) == 3
+        assert await redis_client.zadd(key2, member_scores2) == 1
+        assert await redis_client.zadd(key3, member_scores3) == 4
+
+        assert await redis_client.zdiff([key1, key2]) == ["one", "three"]
+        assert await redis_client.zdiff([key1, key3]) == []
+        assert await redis_client.zdiff([non_existing_key, key3]) == []
+
+        zdiff_map = await redis_client.zdiff_withscores([key1, key2])
+        expected_map = {
+            "one": 1.0,
+            "three": 3.0,
+        }
+        assert compare_maps(zdiff_map, expected_map) is True
+        assert (
+            compare_maps(await redis_client.zdiff_withscores([key1, key3]), {}) is True
+        )
+        assert (
+            compare_maps(
+                await redis_client.zdiff_withscores([non_existing_key, key3]), {}
+            )
+            is True
+        )
+
+        # invalid argument - key list must not be empty
+        with pytest.raises(RequestError):
+            await redis_client.zdiff([])
+
+        # invalid argument - key list must not be empty
+        with pytest.raises(RequestError):
+            await redis_client.zdiff_withscores([])
+
+        # key exists, but it is not a sorted set
+        assert await redis_client.set(string_key, "foo") == OK
+        with pytest.raises(RequestError):
+            await redis_client.zdiff([string_key, key2])
+
+        assert await redis_client.set(string_key, "foo") == OK
+        with pytest.raises(RequestError):
+            await redis_client.zdiff_withscores([string_key, key2])
+
+        # same-slot requirement
+        if isinstance(redis_client, RedisClusterClient):
+            with pytest.raises(RequestError) as e:
+                await redis_client.zdiff(["abc", "zxy", "lkn"])
+            assert "CrossSlot" in str(e)
+
+            with pytest.raises(RequestError) as e:
+                await redis_client.zdiff_withscores(["abc", "zxy", "lkn"])
+            assert "CrossSlot" in str(e)
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_zdiffstore(self, redis_client: TRedisClient):

python/python/tests/test_transaction.py

@@ -46,6 +46,7 @@ async def transaction_test(
     key10 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # hyper log log
     key11 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # streams
     key12 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # geo
+    key13 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # sorted set
 
     value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S")
     value2 = get_random_string(5)
@@ -238,6 +239,13 @@ async def transaction_test(
     transaction.zdiffstore(key8, [key8, key8])
     args.append(0)
 
+    transaction.zadd(key13, {"one": 1.0, "two": 2.0})
+    args.append(2)
+    transaction.zdiff([key13, key8])
+    args.append(["one", "two"])
+    transaction.zdiff_withscores([key13, key8])
+    args.append({"one": 1.0, "two": 2.0})
+
     transaction.pfadd(key10, ["a", "b", "c"])
     args.append(1)