Merge pull request #982 from rhymes/rhymes/redis-62-lpop-rpop-count-argument

[Redis 6.2] Add count argument to lpop and rpop

Author: byroot

lib/redis.rb

@@ -1172,23 +1172,29 @@ def rpushx(key, value)
     end
   end
 
-  # Remove and get the first element in a list.
+  # Remove and get the first elements in a list.
   #
   # @param [String] key
-  # @return [String]
-  def lpop(key)
+  # @param [Integer] count number of elements to remove
+  # @return [String, Array<String>] the values of the first elements
+  def lpop(key, count = nil)
     synchronize do |client|
-      client.call([:lpop, key])
+      command = [:lpop, key]
+      command << count if count
+      client.call(command)
     end
   end
 
-  # Remove and get the last element in a list.
+  # Remove and get the last elements in a list.
   #
   # @param [String] key
-  # @return [String]
-  def rpop(key)
+  # @param [Integer] count number of elements to remove
+  # @return [String, Array<String>] the values of the last elements
+  def rpop(key, count = nil)
     synchronize do |client|
-      client.call([:rpop, key])
+      command = [:rpop, key]
+      command << count if count
+      client.call(command)
     end
   end
 

lib/redis/distributed.rb

@@ -413,14 +413,14 @@ def rpushx(key, value)
       node_for(key).rpushx(key, value)
     end
 
-    # Remove and get the first element in a list.
-    def lpop(key)
-      node_for(key).lpop(key)
+    # Remove and get the first elements in a list.
+    def lpop(key, count = nil)
+      node_for(key).lpop(key, count)
     end
 
-    # Remove and get the last element in a list.
-    def rpop(key)
-      node_for(key).rpop(key)
+    # Remove and get the last elements in a list.
+    def rpop(key, count = nil)
+      node_for(key).rpop(key, count)
     end
 
     # Remove the last element in a list, append it to another list and return

test/lint/lists.rb

@@ -119,6 +119,17 @@ def test_lpop
       assert_equal 1, r.llen("foo")
     end
 
+    def test_lpop_count
+      target_version("6.2") do
+        r.rpush "foo", "s1"
+        r.rpush "foo", "s2"
+
+        assert_equal 2, r.llen("foo")
+        assert_equal ["s1", "s2"], r.lpop("foo", 2)
+        assert_equal 0, r.llen("foo")
+      end
+    end
+
     def test_rpop
       r.rpush "foo", "s1"
       r.rpush "foo", "s2"
@@ -128,6 +139,17 @@ def test_rpop
       assert_equal 1, r.llen("foo")
     end
 
+    def test_rpop_count
+      target_version("6.2") do
+        r.rpush "foo", "s1"
+        r.rpush "foo", "s2"
+
+        assert_equal 2, r.llen("foo")
+        assert_equal ["s2", "s1"], r.rpop("foo", 2)
+        assert_equal 0, r.llen("foo")
+      end
+    end
+
     def test_linsert
       r.rpush "foo", "s1"
       r.rpush "foo", "s3"