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

Co-authored-by: Chayim <[email protected]>

Author: wavenator

docker/base/Dockerfile

@@ -1 +1 @@
-FROM redis:6.2.5-buster
+FROM redis:6.2.5-buster

redis/client.py

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

tests/test_commands.py

@@ -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
@@ -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')