From b80fb2ae708c2a66d84720f69124f5d5c92dc1ef Mon Sep 17 00:00:00 2001
From: Yi-Pin Chen <yi-pin.chen@improving.com>
Date: Wed, 22 May 2024 17:30:26 -0700
Subject: [PATCH] Added Python: SETRANGE command

---
 python/python/glide/async_commands/core.py    | 30 +++++++++++++++++++
 .../glide/async_commands/transaction.py       | 20 +++++++++++++
 python/python/tests/test_async_client.py      | 22 ++++++++++++++
 python/python/tests/test_transaction.py       |  2 ++
 4 files changed, 74 insertions(+)

diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py
index 43466af244..4f18765647 100644
--- a/python/python/glide/async_commands/core.py
+++ b/python/python/glide/async_commands/core.py
@@ -599,6 +599,36 @@ async def incrbyfloat(self, key: str, amount: float) -> float:
             await self._execute_command(RequestType.IncrByFloat, [key, str(amount)]),
         )
 
+    async def setrange(self, key: str, offset: int, value: str) -> int:
+        """
+        Overwrites part of the string stored at `key`, starting at the specified
+        `offset`, for the entire length of `value`.
+        If the `offset` is larger than the current length of the string at `key`,
+        the string is padded with zero bytes to make `offset` fit. Create the `key`
+        if it doesn't exist.
+
+        See https://valkey.io/commands/setrange for more details.
+
+        Args:
+            key (str): The key of the string to update.
+            offset (int): The position in the string where `value` should be written.
+            value (str): The string written with `offset`.
+
+        Returns:
+            int: The length of the string stored at `key` after it was modified.
+
+        Examples:
+            >>> await client.set("key", "Hello World")
+            >>> await client.setrange("key", 6, "Redis")
+                11  # The length of the string stored at `key` after it was modified.
+        """
+        return cast(
+            int,
+            await self._execute_command(
+                RequestType.SetRange, [key, str(offset), value]
+            ),
+        )
+
     async def mset(self, key_value_map: Mapping[str, str]) -> TOK:
         """
         Set multiple keys to multiple values in a single atomic operation.
diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py
index d9b4c3ff5a..9b2c3187bd 100644
--- a/python/python/glide/async_commands/transaction.py
+++ b/python/python/glide/async_commands/transaction.py
@@ -403,6 +403,26 @@ def decrby(self: TTransaction, key: str, amount: int) -> TTransaction:
         """
         return self.append_command(RequestType.DecrBy, [key, str(amount)])
 
+    def setrange(self: TTransaction, key: str, offset: int, value: str) -> TTransaction:
+        """
+        Overwrites part of the string stored at `key`, starting at the specified
+        `offset`, for the entire length of `value`.
+        If the `offset` is larger than the current length of the string at `key`,
+        the string is padded with zero bytes to make `offset` fit. Create the `key`
+        if it doesn't exist.
+
+        See https://valkey.io/commands/setrange for more details.
+
+        Args:
+            key (str): The key of the string to update.
+            offset (int): The position in the string where `value` should be written.
+            value (str): The string written with `offset`.
+
+        Command response:
+            int: The length of the string stored at `key` after it was modified.
+        """
+        return self.append_command(RequestType.SetRange, [key, str(offset), value])
+
     def hset(
         self: TTransaction, key: str, field_value_map: Mapping[str, str]
     ) -> TTransaction:
diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py
index 095ec654e8..ff27b8e7f3 100644
--- a/python/python/tests/test_async_client.py
+++ b/python/python/tests/test_async_client.py
@@ -562,6 +562,28 @@ async def test_decr_with_str_value(self, redis_client: TRedisClient):
 
         assert "value is not an integer" in str(e)
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_setrange(self, redis_client: TRedisClient):
+        key1 = get_random_string(10)
+        key2 = get_random_string(10)
+
+        # test new key and existing key
+        assert await redis_client.setrange(key1, 0, "Hello World") == 11
+        assert await redis_client.setrange(key1, 6, "GLIDE") == 11
+
+        # offset > len
+        assert await redis_client.setrange(key1, 15, "GLIDE") == 20
+
+        # negative offset
+        with pytest.raises(RequestError):
+            assert await redis_client.setrange(key1, -1, "GLIDE")
+
+        # non-string key
+        assert await redis_client.lpush(key2, ["_"]) == 1
+        with pytest.raises(RequestError):
+            assert await redis_client.setrange(key2, 0, "_")
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_hset_hget_hgetall(self, redis_client: TRedisClient):
diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py
index 5856aef735..30d4306a36 100644
--- a/python/python/tests/test_transaction.py
+++ b/python/python/tests/test_transaction.py
@@ -63,6 +63,8 @@ async def transaction_test(
 
     transaction.set(key, value)
     args.append(OK)
+    transaction.setrange(key, 0, value)
+    args.append(len(value))
     transaction.get(key)
     args.append(value)
     transaction.type(key)