Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix issue #44 - Add 'score_cast_func' parameters to sorted sets #45

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions fakeredis.py
Original file line number Diff line number Diff line change
Expand Up @@ -828,7 +828,7 @@ def zinterstore(self, dest, keys, aggregate=None):
lambda x: x in
valid_keys)

def zrange(self, name, start, end, desc=False, withscores=False):
def zrange(self, name, start, end, desc=False, withscores=False, score_cast_func=float):
"""
Return a range of values from sorted set ``name`` between
``start`` and ``end`` sorted in ascending order.
Expand All @@ -854,15 +854,15 @@ def zrange(self, name, start, end, desc=False, withscores=False):
if not withscores:
return items
else:
return [(k, all_items[k]) for k in items]
return [(k, score_cast_func(all_items[k])) for k in items]

def _get_zelements_in_order(self, all_items, reverse=False):
by_keyname = sorted(all_items.items(), key=lambda x: x[0])
in_order = sorted(by_keyname, key=lambda x: x[1], reverse=reverse)
return [el[0] for el in in_order]

def zrangebyscore(self, name, min, max,
start=None, num=None, withscores=False):
start=None, num=None, withscores=False, score_cast_func=float):
"""
Return a range of values from the sorted set ``name`` with scores
between ``min`` and ``max``.
Expand All @@ -874,9 +874,9 @@ def zrangebyscore(self, name, min, max,
The return type is a list of (value, score) pairs
"""
return self._zrangebyscore(name, min, max, start, num, withscores,
reverse=False)
reverse=False, score_cast_func=score_cast_func)

def _zrangebyscore(self, name, min, max, start, num, withscores, reverse):
def _zrangebyscore(self, name, min, max, start, num, withscores, reverse, score_cast_func=float):
if (start is not None and num is None) or \
(num is not None and start is None):
raise redis.RedisError("``start`` and ``num`` must both "
Expand All @@ -890,7 +890,7 @@ def _zrangebyscore(self, name, min, max, start, num, withscores, reverse):
if start is not None:
matches = matches[start:start + num]
if withscores:
return [(k, all_items[k]) for k in matches]
return [(k, score_cast_func(all_items[k])) for k in matches]
return matches

def zrank(self, name, value):
Expand Down Expand Up @@ -947,7 +947,7 @@ def zremrangebyscore(self, name, min, max):
removed += 1
return removed

def zrevrange(self, name, start, num, withscores=False):
def zrevrange(self, name, start, num, withscores=False, score_cast_func=float):
"""
Return a range of values from sorted set ``name`` between
``start`` and ``num`` sorted in descending order.
Expand All @@ -957,10 +957,10 @@ def zrevrange(self, name, start, num, withscores=False):
``withscores`` indicates to return the scores along with the values
The return type is a list of (value, score) pairs
"""
return self.zrange(name, start, num, True, withscores)
return self.zrange(name, start, num, True, withscores, score_cast_func)

def zrevrangebyscore(self, name, max, min,
start=None, num=None, withscores=False):
start=None, num=None, withscores=False, score_cast_func=float):
"""
Return a range of values from the sorted set ``name`` with scores
between ``min`` and ``max`` in descending order.
Expand All @@ -972,7 +972,7 @@ def zrevrangebyscore(self, name, max, min,
The return type is a list of (value, score) pairs
"""
return self._zrangebyscore(name, min, max, start, num, withscores,
reverse=True)
reverse=True, score_cast_func=score_cast_func)

def zrevrank(self, name, value):
"""
Expand Down
15 changes: 15 additions & 0 deletions test_fakeredis.py
Original file line number Diff line number Diff line change
Expand Up @@ -867,6 +867,14 @@ def test_zrange_descending_with_scores(self):
withscores=True),
[('three', 3), ('two', 2), ('one', 1)])

def test_zrange_descending_with_scores_cast_func(self):
self.redis.zadd('foo', one=1)
self.redis.zadd('foo', two=2)
self.redis.zadd('foo', three=3)
self.assertEqual(self.redis.zrange('foo', 0, -1, desc=True,
withscores=True, score_cast_func=str),
[('three', "3.0"), ('two', "2.0"), ('one', "1.0")])

def test_zrange_with_positive_indices(self):
self.redis.zadd('foo', one=1)
self.redis.zadd('foo', two=2)
Expand Down Expand Up @@ -958,6 +966,13 @@ def test_zrangebyscore_withscores(self):
self.assertEqual(self.redis.zrangebyscore('foo', 1, 3, 0, 2, True),
[('one', 1), ('two', 2)])

def test_zrangebyscore_withscores_cast_func(self):
self.redis.zadd('foo', one=1)
self.redis.zadd('foo', two=2)
self.redis.zadd('foo', three=3)
self.assertEqual(self.redis.zrangebyscore('foo', 1, 3, 0, 2, True, score_cast_func=str),
[('one', "1.0"), ('two', "2.0")])

def test_zrevrangebyscore(self):
self.redis.zadd('foo', one=1)
self.redis.zadd('foo', two=2)
Expand Down