Python: adds ZREM command and ZADD for transaction

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),
+        )

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):
     """

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):

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,
     ]