Skip to content

Commit

Permalink
Add a count parameter to lpop/rpop for redis >= 6.2.0 (#1487)
Browse files Browse the repository at this point in the history
Co-authored-by: Chayim <chayim@users.noreply.github.com>
  • Loading branch information
wavenator and chayim authored Aug 5, 2021
1 parent 5240d60 commit 238f69e
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 7 deletions.
2 changes: 1 addition & 1 deletion docker/base/Dockerfile
Original file line number Diff line number Diff line change
@@ -1 +1 @@
FROM redis:6.2.5-buster
FROM redis:6.2.5-buster
30 changes: 24 additions & 6 deletions redis/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2148,9 +2148,18 @@ def llen(self, name):
"Return the length of the list ``name``"
return self.execute_command('LLEN', name)

def lpop(self, name):
"Remove and return the first item of the list ``name``"
return self.execute_command('LPOP', name)
def lpop(self, name, count=None):
"""
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, *values):
"Push ``values`` onto the head of the list ``name``"
Expand Down Expand Up @@ -2196,9 +2205,18 @@ def ltrim(self, name, start, end):
"""
return self.execute_command('LTRIM', name, start, end)

def rpop(self, name):
"Remove and return the last item of the list ``name``"
return self.execute_command('RPOP', name)
def rpop(self, name, count=None):
"""
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, dst):
"""
Expand Down
16 changes: 16 additions & 0 deletions tests/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -1122,6 +1122,14 @@ def test_lpop(self, r):
assert r.lpop('a') == b'3'
assert r.lpop('a') is None

@skip_if_server_version_lt('6.2.0')
def test_lpop_count(self, r):
r.rpush('a', '1', '2', '3')
assert r.lpop('a', 2) == [b'1', b'2']
assert r.lpop('a', 1) == [b'3']
assert r.lpop('a') is None
assert r.lpop('a', 3) is None

def test_lpush(self, r):
assert r.lpush('a', '1') == 1
assert r.lpush('a', '2') == 2
Expand Down Expand Up @@ -1171,6 +1179,14 @@ def test_rpop(self, r):
assert r.rpop('a') == b'1'
assert r.rpop('a') is None

@skip_if_server_version_lt('6.2.0')
def test_rpop_count(self, r):
r.rpush('a', '1', '2', '3')
assert r.rpop('a', 2) == [b'3', b'2']
assert r.rpop('a', 1) == [b'1']
assert r.rpop('a') is None
assert r.rpop('a', 3) is None

def test_rpoplpush(self, r):
r.rpush('a', 'a1', 'a2', 'a3')
r.rpush('b', 'b1', 'b2', 'b3')
Expand Down

0 comments on commit 238f69e

Please sign in to comment.