From 961da8cc1ba04c3fbb3de0485bd7b89bdbf97c03 Mon Sep 17 00:00:00 2001 From: Andrew-Chen-Wang Date: Tue, 5 Oct 2021 10:52:49 -0400 Subject: [PATCH] Add andymccurdy/redis-py@238f69e36e0ff5ac9b892706e3a5478b138069cd Add a count parameter to lpop/rpop for redis >= 6.2.0 (andymccurdy/redis-py#1487) Signed-off-by: Andrew-Chen-Wang --- aioredis/client.py | 30 ++++++++++++++++++++++++------ tests/test_commands.py | 16 ++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/aioredis/client.py b/aioredis/client.py index d32e67851..5c52478ac 100644 --- a/aioredis/client.py +++ b/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: """ diff --git a/tests/test_commands.py b/tests/test_commands.py index b15a4cf6f..fe3fae14f 100644 --- a/tests/test_commands.py +++ b/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")