From fb709081368d76202a7cc113c7ae20536e159faa Mon Sep 17 00:00:00 2001
From: Hugo Hache <hugo.hache@sorare.com>
Date: Mon, 9 Dec 2024 17:45:16 +0100
Subject: [PATCH] Add :with_score option to zrank and zrevrank

---
 lib/redis/commands.rb             |  8 ++++--
 lib/redis/commands/sorted_sets.rb | 46 +++++++++++++++++++++++++++----
 lib/redis/distributed.rb          |  8 +++---
 test/lint/sorted_sets.rb          |  8 ++++++
 4 files changed, 57 insertions(+), 13 deletions(-)

diff --git a/lib/redis/commands.rb b/lib/redis/commands.rb
index 179c2200a..da7354bc1 100644
--- a/lib/redis/commands.rb
+++ b/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|
diff --git a/lib/redis/commands/sorted_sets.rb b/lib/redis/commands/sorted_sets.rb
index 4b60b4739..95c60a76f 100644
--- a/lib/redis/commands/sorted_sets.rb
+++ b/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.
diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb
index fb156fbab..eef03803b 100644
--- a/lib/redis/distributed.rb
+++ b/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.
diff --git a/test/lint/sorted_sets.rb b/test/lint/sorted_sets.rb
index 10f8ce5e9..f562af130 100644
--- a/test/lint/sorted_sets.rb
+++ b/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