Add redis/redis-py@238f69e

Add a count parameter to lpop/rpop for redis >= 6.2.0 (redis/redis-py#1487)

Signed-off-by: Andrew-Chen-Wang <acwangpython@gmail.com>

aioredis/client.py

@@ -2318,9 +2318,18 @@ def llen(self, name: KeyT) -> Awaitable:
         """Return the length of the list ``name``"""
         return self.execute_command("LLEN", name)
 
-    def lpop(self, name: KeyT) -> Awaitable:
-        """Remove and return the first item of the list ``name``"""
-        return self.execute_command("LPOP", name)
+    def lpop(self, name: KeyT, count: Optional[int] = None) -> Awaitable:
+        """
+        Removes and returns the first elements of the list ``name``.
+
+        By default, the command pops a single element from the beginning of
+        the list. When provided with the optional ``count`` argument, the reply
+        will consist of up to count elements, depending on the list's length.
+        """
+        if count is not None:
+            return self.execute_command("LPOP", name, count)
+        else:
+            return self.execute_command("LPOP", name)
 
     def lpush(self, name: KeyT, *values: EncodableT) -> Awaitable:
         """Push ``values`` onto the head of the list ``name``"""
@@ -2366,9 +2375,18 @@ def ltrim(self, name: KeyT, start: int, end: int) -> Awaitable:
         """
         return self.execute_command("LTRIM", name, start, end)
 
-    def rpop(self, name: KeyT) -> Awaitable:
-        """Remove and return the last item of the list ``name``"""
-        return self.execute_command("RPOP", name)
+    def rpop(self, name: KeyT, count: Optional[int] = None) -> Awaitable:
+        """
+        Removes and returns the last elements of the list ``name``.
+
+        By default, the command pops a single element from the end of the list.
+        When provided with the optional ``count`` argument, the reply will
+        consist of up to count elements, depending on the list's length.
+        """
+        if count is not None:
+            return self.execute_command("RPOP", name, count)
+        else:
+            return self.execute_command("RPOP", name)
 
     def rpoplpush(self, src: KeyT, dst: KeyT) -> Awaitable:
         """

tests/test_commands.py

@@ -1190,6 +1190,14 @@ async def test_lpop(self, r: aioredis.Redis):
         assert await r.lpop("a") == b"3"
         assert await r.lpop("a") is None
 
+    @skip_if_server_version_lt("6.2.0")
+    async def test_lpop_count(self, r: aioredis.Redis):
+        await r.rpush('a', '1', '2', '3')
+        assert await r.lpop('a', 2) == [b'1', b'2']
+        assert await r.lpop('a', 1) == [b'3']
+        assert await r.lpop('a') is None
+        assert await r.lpop('a', 3) is None
+
     async def test_lpush(self, r: aioredis.Redis):
         assert await r.lpush("a", "1") == 1
         assert await r.lpush("a", "2") == 2
@@ -1239,6 +1247,14 @@ async def test_rpop(self, r: aioredis.Redis):
         assert await r.rpop("a") == b"1"
         assert await r.rpop("a") is None
 
+    @skip_if_server_version_lt('6.2.0')
+    async def test_rpop_count(self, r: aioredis.Redis):
+        await r.rpush('a', '1', '2', '3')
+        assert await r.rpop('a', 2) == [b'3', b'2']
+        assert await r.rpop('a', 1) == [b'1']
+        assert await r.rpop('a') is None
+        assert await r.rpop('a', 3) is None
+
     async def test_rpoplpush(self, r: aioredis.Redis):
         await r.rpush("a", "a1", "a2", "a3")
         await r.rpush("b", "b1", "b2", "b3")