Skip to content

Commit

Permalink
Add :with_score option to zrank and zrevrank
Browse files Browse the repository at this point in the history
  • Loading branch information
Hugo Hache committed Dec 9, 2024
1 parent 08f3e18 commit 2e7c468
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 13 deletions.
8 changes: 5 additions & 3 deletions lib/redis/commands.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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|
Expand Down
46 changes: 40 additions & 6 deletions lib/redis/commands/sorted_sets.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 4 additions & 4 deletions lib/redis/distributed.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions test/lint/sorted_sets.rb
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,8 @@ def test_zrank
r.zadd "foo", 3, "s3"

assert_equal 2, r.zrank("foo", "s3")
assert_equal [2, 3], r.zrank("foo", "s3", with_score: true)
assert_equal [2, 3], r.zrank("foo", "s3", withscore: true)
end

def test_zrevrank
Expand All @@ -234,6 +236,8 @@ def test_zrevrank
r.zadd "foo", 3, "s3"

assert_equal 0, r.zrevrank("foo", "s3")
assert_equal [0, 3], r.zrevrank("foo", "s3", with_score: true)
assert_equal [0, 3], r.zrevrank("foo", "s3", withscore: true)
end

def test_zrange
Expand Down

0 comments on commit 2e7c468

Please sign in to comment.