Skip to content

Commit d95d8a2

Browse files
authored
Add WITHSCORE to ZRANK (redis#2758)
* add withscore to zrank with tests * fix test
1 parent db7b9dd commit d95d8a2

File tree

3 files changed

+29
-2
lines changed

3 files changed

+29
-2
lines changed

redis/commands/core.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4654,13 +4654,22 @@ def zrevrangebyscore(
46544654
options = {"withscores": withscores, "score_cast_func": score_cast_func}
46554655
return self.execute_command(*pieces, **options)
46564656

4657-
def zrank(self, name: KeyT, value: EncodableT) -> ResponseT:
4657+
def zrank(
4658+
self,
4659+
name: KeyT,
4660+
value: EncodableT,
4661+
withscore: bool = False,
4662+
) -> ResponseT:
46584663
"""
46594664
Returns a 0-based value indicating the rank of ``value`` in sorted set
4660-
``name``
4665+
``name``.
4666+
The optional WITHSCORE argument supplements the command's
4667+
reply with the score of the element returned.
46614668
46624669
For more information see https://redis.io/commands/zrank
46634670
"""
4671+
if withscore:
4672+
return self.execute_command("ZRANK", name, value, "WITHSCORE")
46644673
return self.execute_command("ZRANK", name, value)
46654674

46664675
def zrem(self, name: KeyT, *values: FieldT) -> ResponseT:

tests/test_asyncio/test_commands.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1645,6 +1645,15 @@ async def test_zrank(self, r: redis.Redis):
16451645
assert await r.zrank("a", "a2") == 1
16461646
assert await r.zrank("a", "a6") is None
16471647

1648+
@skip_if_server_version_lt("7.2.0")
1649+
async def test_zrank_withscore(self, r: redis.Redis):
1650+
await r.zadd("a", {"a1": 1, "a2": 2, "a3": 3, "a4": 4, "a5": 5})
1651+
assert await r.zrank("a", "a1") == 0
1652+
assert await r.rank("a", "a2") == 1
1653+
assert await r.zrank("a", "a6") is None
1654+
assert await r.zrank("a", "a3", withscore=True) == [2, "3"]
1655+
assert await r.zrank("a", "a6", withscore=True) is None
1656+
16481657
async def test_zrem(self, r: redis.Redis):
16491658
await r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
16501659
assert await r.zrem("a", "a2") == 1

tests/test_commands.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2581,6 +2581,15 @@ def test_zrank(self, r):
25812581
assert r.zrank("a", "a2") == 1
25822582
assert r.zrank("a", "a6") is None
25832583

2584+
@skip_if_server_version_lt("7.2.0")
2585+
def test_zrank_withscore(self, r: redis.Redis):
2586+
r.zadd("a", {"a1": 1, "a2": 2, "a3": 3, "a4": 4, "a5": 5})
2587+
assert r.zrank("a", "a1") == 0
2588+
assert r.rank("a", "a2") == 1
2589+
assert r.zrank("a", "a6") is None
2590+
assert r.zrank("a", "a3", withscore=True) == [2, "3"]
2591+
assert r.zrank("a", "a6", withscore=True) is None
2592+
25842593
def test_zrem(self, r):
25852594
r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
25862595
assert r.zrem("a", "a2") == 1

0 commit comments

Comments
 (0)