From 3af09e0499c9ab7261efd205b96b7bc1b737657f Mon Sep 17 00:00:00 2001 From: Aaron <69273634+aaron-congo@users.noreply.github.com> Date: Mon, 13 May 2024 16:31:21 -0700 Subject: [PATCH] Python: add ZDIFF command (#1401) * Python: add ZDIFF command (#255) * Add PR link in CHANGELOG * PR suggestions --- CHANGELOG.md | 2 +- python/python/glide/async_commands/core.py | 59 +++++++++++++++++ .../glide/async_commands/transaction.py | 35 ++++++++++ python/python/tests/test_async_client.py | 64 +++++++++++++++++++ python/python/tests/test_transaction.py | 8 +++ 5 files changed, 167 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b38788bf15..41e5270499 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 921535ab84..4267bd7c63 100644 --- a/python/python/glide/async_commands/core.py +++ b/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 diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index dbf3651d8f..43ad39e59d 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/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: diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 0461e8132a..7d87aaa382 100644 --- a/python/python/tests/test_async_client.py +++ b/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): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index df45b3b8e4..83bac56b35 100644 --- a/python/python/tests/test_transaction.py +++ b/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)