From 04cde134017bfba7c56413614df0ad732713483c Mon Sep 17 00:00:00 2001 From: Shoham Elias Date: Thu, 18 Jan 2024 11:35:22 +0000 Subject: [PATCH] Python: adds ZREM command --- python/python/glide/async_commands/core.py | 31 +++++ .../glide/async_commands/transaction.py | 125 ++++++++++++++++++ python/python/tests/test_async_client.py | 11 ++ python/python/tests/test_transaction.py | 7 + 4 files changed, 174 insertions(+) diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index c231f7a4b9..ff1f0a5fa5 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -1185,3 +1185,34 @@ async def zadd_incr( Optional[float], await self._execute_command(RequestType.Zadd, args), ) + + async def zrem( + self, + key: str, + members: List[str], + ) -> int: + """ + Removes the specified members from the sorted set stored at `key`. + Specified members that are not a member of this set are ignored. + + See https://redis.io/commands/zrem/ for more details. + + Args: + key (str): The key of the sorted set. + members (List[str]): A list of members to remove from the sorted set. + + Returns: + int: The number of members that were removed from the sorted set, not including non-existing members. + If `key` does not exist, it is treated as an empty sorted set, and this command returns 0. + If `key` holds a value that is not a sorted set, an error is returned. + + Examples: + >>> await zrem("my_sorted_set", ["member1", "member2"]) + 2 # Indicates that two members have been removed from the sorted set "my_sorted_set." + >>> await zrem("non_existing_sorted_set", ["member1", "member2"]) + 0 # Indicates that no members were removed as the sorted set "non_existing_sorted_set" does not exist. + """ + return cast( + int, + await self._execute_command(RequestType.Zrem, [key] + members), + ) diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index 67b31c026c..47b757cef9 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -6,6 +6,7 @@ ExpireOptions, ExpirySet, InfoSection, + UpdateOptions, ) from glide.protobuf.redis_request_pb2 import RequestType @@ -774,6 +775,130 @@ def ttl(self, key: str): """ self.append_command(RequestType.TTL, [key]) + def zadd( + self, + key: str, + members_scores: Mapping[str, float], + existing_options: Optional[ConditionalChange] = None, + update_condition: Optional[UpdateOptions] = None, + changed: bool = False, + ): + """ + Adds members with their scores to the sorted set stored at `key`. + If a member is already a part of the sorted set, its score is updated. + + See https://redis.io/commands/zadd/ for more details. + + Args: + key (str): The key of the sorted set. + members_scores (Mapping[str, float]): A mapping of members to their corresponding scores. + existing_options (Optional[ConditionalChange]): Options for handling existing members. + - NX: Only add new elements. + - XX: Only update existing elements. + update_condition (Optional[UpdateOptions]): Options for updating scores. + - GT: Only update scores greater than the current values. + - LT: Only update scores less than the current values. + changed (bool): Modify the return value to return the number of changed elements, instead of the number of new elements added. + + Commands response: + int: The number of elements added to the sorted set. + If `changed` is set, returns the number of elements updated in the sorted set. + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if update_condition: + args.append(update_condition.value) + + if changed: + args.append("CH") + + if existing_options and update_condition: + if existing_options == ConditionalChange.ONLY_IF_DOES_NOT_EXIST: + raise ValueError( + "The GT, LT and NX options are mutually exclusive. " + f"Cannot choose both {update_condition.value} and NX." + ) + + members_scores_list = [ + str(item) for pair in members_scores.items() for item in pair[::-1] + ] + args += members_scores_list + + self.append_command(RequestType.Zadd, args) + + def zadd_incr( + self, + key: str, + member: str, + increment: float, + existing_options: Optional[ConditionalChange] = None, + update_condition: Optional[UpdateOptions] = None, + ): + """ + Increments the score of member in the sorted set stored at `key` by `increment`. + If `member` does not exist in the sorted set, it is added with `increment` as its score (as if its previous score was 0.0). + If `key` does not exist, a new sorted set with the specified member as its sole member is created. + + See https://redis.io/commands/zadd/ for more details. + + Args: + key (str): The key of the sorted set. + member (str): A member in the sorted set to increment. + increment (float): The score to increment the member. + existing_options (Optional[ConditionalChange]): Options for handling the member's existence. + - NX: Only increment a member that doesn't exist. + - XX: Only increment an existing member. + update_condition (Optional[UpdateOptions]): Options for updating the score. + - GT: Only increment the score of the member if the new score will be greater than the current score. + - LT: Only increment (decrement) the score of the member if the new score will be less than the current score. + + Commands response: + Optional[float]: The score of the member. + If there was a conflict with choosing the XX/NX/LT/GT options, the operation aborts and null is returned. + """ + args = [key] + if existing_options: + args.append(existing_options.value) + + if update_condition: + args.append(update_condition.value) + + args.append("INCR") + + if existing_options and update_condition: + if existing_options == ConditionalChange.ONLY_IF_DOES_NOT_EXIST: + raise ValueError( + "The GT, LT and NX options are mutually exclusive. " + f"Cannot choose both {update_condition.value} and NX." + ) + + args += [str(increment), member] + self.append_command(RequestType.Zadd, args) + + def zrem( + self, + key: str, + members: List[str], + ): + """ + Removes the specified members from the sorted set stored at `key`. + Specified members that are not a member of this set are ignored. + + See https://redis.io/commands/zrem/ for more details. + + Args: + key (str): The key of the sorted set. + members (List[str]): A list of members to remove from the sorted set. + + Commands response: + int: The number of members that were removed from the sorted set, not including non-existing members. + If `key` does not exist, it is treated as an empty sorted set, and this command returns 0. + If `key` holds a value that is not a sorted set, an error is returned. + """ + self.append_command(RequestType.Zrem, [key] + members) + class Transaction(BaseTransaction): """ diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 83f54870d4..23052d7b1d 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -1044,6 +1044,17 @@ async def test_zadd_gt_lt(self, redis_client: TRedisClient): == None ) + @pytest.mark.parametrize("cluster_mode", [True, False]) + async def test_zrem(self, redis_client: TRedisClient): + key = get_random_string(10) + members_scores = {"one": 1, "two": 2, "three": 3} + assert await redis_client.zadd(key, members_scores=members_scores) == 3 + + assert await redis_client.zrem(key, ["one"]) == 1 + assert await redis_client.zrem(key, ["one", "two", "three"]) == 2 + + assert await redis_client.zrem("non_existing_set", ["member"]) == 0 + class TestCommandsUnitTests: def test_expiry_cmd_args(self): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 6f662964c9..20e1c783aa 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -25,6 +25,7 @@ def transaction_test( key5 = "{{{}}}:{}".format(keyslot, get_random_string(3)) key6 = "{{{}}}:{}".format(keyslot, get_random_string(3)) key7 = "{{{}}}:{}".format(keyslot, get_random_string(3)) + key8 = "{{{}}}:{}".format(keyslot, get_random_string(3)) value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S") value2 = get_random_string(5) @@ -83,6 +84,9 @@ def transaction_test( transaction.smembers(key7) transaction.scard(key7) + transaction.zadd(key8, {"one": 1, "two": 2, "three": 3}) + transaction.zadd_incr(key8, "one", 3) + transaction.zrem(key8, ["one"]) return [ OK, value, @@ -121,6 +125,9 @@ def transaction_test( 1, {"bar"}, 1, + 3, + 4, + 1, ]