Skip to content

Commit

Permalink
Python: adds ZREM command
Browse files Browse the repository at this point in the history
  • Loading branch information
shohamazon committed Jan 18, 2024
1 parent 9d90474 commit 04cde13
Show file tree
Hide file tree
Showing 4 changed files with 174 additions and 0 deletions.
31 changes: 31 additions & 0 deletions python/python/glide/async_commands/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)
125 changes: 125 additions & 0 deletions python/python/glide/async_commands/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
ExpireOptions,
ExpirySet,
InfoSection,
UpdateOptions,
)
from glide.protobuf.redis_request_pb2 import RequestType

Expand Down Expand Up @@ -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):
"""
Expand Down
11 changes: 11 additions & 0 deletions python/python/tests/test_async_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
7 changes: 7 additions & 0 deletions python/python/tests/test_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -121,6 +125,9 @@ def transaction_test(
1,
{"bar"},
1,
3,
4,
1,
]


Expand Down

0 comments on commit 04cde13

Please sign in to comment.