Python: add ZDIFFSTORE command (valkey-io#1378)

* Python: add ZDIFFSTORE command (#257)

* Update CHANGELOG with PR link

* PR suggestions

* PR suggestions

* PR suggestions

---------

Co-authored-by: Andrew Carbonetto <andrew.carbonetto@improving.com>

CHANGELOG.md

@@ -23,6 +23,7 @@
 * Python: Added XADD, XTRIM commands ([#1320](https://github.com/aws/glide-for-redis/pull/1320))
 * 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))
 
 
 #### Fixes

python/python/glide/async_commands/core.py

@@ -2720,6 +2720,40 @@ async def zmscore(
             await self._execute_command(RequestType.ZMScore, [key] + members),
         )
 
+    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
+        the difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are
+        treated as empty sets.
+
+        When in Cluster mode, all keys in `keys` and `destination` must map to the same hash slot.
+
+        See https://valkey.io/commands/zdiffstore for more details.
+
+        Args:
+            destination (str): The key for the resulting sorted set.
+            keys (List[str]): The keys of the sorted sets to compare.
+
+        Returns:
+            int: The number of members in the resulting sorted set stored at `destination`.
+
+        Examples:
+            >>> await client.zadd("key1", {"member1": 10.5, "member2": 8.2})
+                2  # Indicates that two elements have been added to the sorted set at "key1".
+            >>> await client.zadd("key2", {"member1": 10.5})
+                1  # Indicates that one element has been added to the sorted set at "key2".
+            >>> await client.zdiffstore("my_sorted_set", ["key1", "key2"])
+                1  # One member exists in "key1" but not "key2", and this member was stored in "my_sorted_set".
+            >>> await client.zrange("my_sorted_set", RangeByIndex(0, -1))
+                ['member2']  # "member2" is now stored in "my_sorted_set"
+        """
+        return cast(
+            int,
+            await self._execute_command(
+                RequestType.ZDiffStore, [destination, str(len(keys))] + keys
+            ),
+        )
+
     async def invoke_script(
         self,
         script: Script,

python/python/glide/async_commands/transaction.py

@@ -1947,6 +1947,27 @@ def zmscore(self: TTransaction, key: str, members: List[str]) -> TTransaction:
         """
         return self.append_command(RequestType.ZMScore, [key] + members)
 
+    def zdiffstore(
+        self: TTransaction, destination: str, keys: List[str]
+    ) -> TTransaction:
+        """
+        Calculates the difference between the first sorted set and all the successive sorted sets at `keys` and stores
+        the difference as a sorted set to `destination`, overwriting it if it already exists. Non-existent keys are
+        treated as empty sets.
+
+        See https://valkey.io/commands/zdiffstore for more details.
+
+        Args:
+            destination (str): The key for the resulting sorted set.
+            keys (List[str]): The keys of the sorted sets to compare.
+
+        Command response:
+            int: The number of members in the resulting sorted set stored at `destination`.
+        """
+        return self.append_command(
+            RequestType.ZDiffStore, [destination, str(len(keys))] + keys
+        )
+
     def dbsize(self: TTransaction) -> TTransaction:
         """
         Returns the number of keys in the currently selected database.

python/python/tests/test_async_client.py

@@ -2380,6 +2380,52 @@ 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_zdiffstore(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)}"
+        key4 = f"{{testKey}}:4-{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.zdiffstore(key4, [key1, key2]) == 2
+        assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {
+            "one": 1.0,
+            "three": 3.0,
+        }
+
+        assert await redis_client.zdiffstore(key4, [key3, key2, key1]) == 1
+        assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {
+            "four": 4.0
+        }
+
+        assert await redis_client.zdiffstore(key4, [key1, key3]) == 0
+        assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {}
+
+        assert await redis_client.zdiffstore(key4, [non_existing_key, key1]) == 0
+        assert await redis_client.zrange_withscores(key4, RangeByIndex(0, -1)) == {}
+
+        # key exists, but it is not a sorted set
+        assert await redis_client.set(string_key, "value") == OK
+        with pytest.raises(RequestError):
+            await redis_client.zdiffstore(key4, [string_key, key1])
+
+        # same-slot requirement
+        if isinstance(redis_client, RedisClusterClient):
+            with pytest.raises(RequestError) as e:
+                await redis_client.zdiffstore("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_type(self, redis_client: TRedisClient):

python/python/tests/test_transaction.py

@@ -235,6 +235,8 @@ async def transaction_test(
     args.append(1)
     transaction.zremrangebylex(key8, InfBound.NEG_INF, InfBound.POS_INF)
     args.append(0)
+    transaction.zdiffstore(key8, [key8, key8])
+    args.append(0)
 
     transaction.pfadd(key10, ["a", "b", "c"])
     args.append(1)