Merge pull request #1300 from Hugo-Hache/hh/zrevrank-withscore

Add :with_score option to zrank and zrevrank

Author: byroot

lib/redis/commands.rb

@@ -83,12 +83,14 @@ module Commands
       end
     }
 
+    FloatifyPair = lambda { |(first, score)|
+      [first, Floatify.call(score)]
+    }
+
     FloatifyPairs = lambda { |value|
       return value unless value.respond_to?(:each_slice)
 
-      value.each_slice(2).map do |member, score|
-        [member, Floatify.call(score)]
-      end
+      value.each_slice(2).map(&FloatifyPair)
     }
 
     HashifyInfo = lambda { |reply|

lib/redis/commands/sorted_sets.rb

@@ -454,21 +454,55 @@ def zrevrange(key, start, stop, withscores: false, with_scores: withscores)
 
       # Determine the index of a member in a sorted set.
       #
+      # @example Retrieve member rank
+      #   redis.zrank("zset", "a")
+      #     # => 3
+      # @example Retrieve member rank with their score
+      #   redis.zrank("zset", "a", :with_score => true)
+      #     # => [3, 32.0]
+      #
       # @param [String] key
       # @param [String] member
-      # @return [Integer]
-      def zrank(key, member)
-        send_command([:zrank, key, member])
+      #
+      # @return [Integer, [Integer, Float]]
+      #   - when `:with_score` is not specified, an Integer
+      #   - when `:with_score` is specified, a `[rank, score]` pair
+      def zrank(key, member, withscore: false, with_score: withscore)
+        args = [:zrank, key, member]
+
+        if with_score
+          args << "WITHSCORE"
+          block = FloatifyPair
+        end
+
+        send_command(args, &block)
       end
 
       # Determine the index of a member in a sorted set, with scores ordered from
       # high to low.
       #
+      # @example Retrieve member rank
+      #   redis.zrevrank("zset", "a")
+      #     # => 3
+      # @example Retrieve member rank with their score
+      #   redis.zrevrank("zset", "a", :with_score => true)
+      #     # => [3, 32.0]
+      #
       # @param [String] key
       # @param [String] member
-      # @return [Integer]
-      def zrevrank(key, member)
-        send_command([:zrevrank, key, member])
+      #
+      # @return [Integer, [Integer, Float]]
+      #   - when `:with_score` is not specified, an Integer
+      #   - when `:with_score` is specified, a `[rank, score]` pair
+      def zrevrank(key, member, withscore: false, with_score: withscore)
+        args = [:zrevrank, key, member]
+
+        if with_score
+          args << "WITHSCORE"
+          block = FloatifyPair
+        end
+
+        send_command(args, &block)
       end
 
       # Remove all members in a sorted set within the given indexes.

lib/redis/distributed.rb

@@ -752,14 +752,14 @@ def zrevrange(key, start, stop, **options)
     end
 
     # Determine the index of a member in a sorted set.
-    def zrank(key, member)
-      node_for(key).zrank(key, member)
+    def zrank(key, member, **options)
+      node_for(key).zrank(key, member, **options)
     end
 
     # Determine the index of a member in a sorted set, with scores ordered from
     # high to low.
-    def zrevrank(key, member)
-      node_for(key).zrevrank(key, member)
+    def zrevrank(key, member, **options)
+      node_for(key).zrevrank(key, member, **options)
     end
 
     # Remove all members in a sorted set within the given indexes.

test/lint/sorted_sets.rb

@@ -226,6 +226,10 @@ def test_zrank
       r.zadd "foo", 3, "s3"
 
       assert_equal 2, r.zrank("foo", "s3")
+      target_version "7.2" do
+        assert_equal [2, 3], r.zrank("foo", "s3", with_score: true)
+        assert_equal [2, 3], r.zrank("foo", "s3", withscore: true)
+      end
     end
 
     def test_zrevrank
@@ -234,6 +238,10 @@ def test_zrevrank
       r.zadd "foo", 3, "s3"
 
       assert_equal 0, r.zrevrank("foo", "s3")
+      target_version "7.2" do
+        assert_equal [0, 3], r.zrevrank("foo", "s3", with_score: true)
+        assert_equal [0, 3], r.zrevrank("foo", "s3", withscore: true)
+      end
     end
 
     def test_zrange