Skip to content

Commit

Permalink
Supports argument score_cast_func in zrangebyscore()
Browse files Browse the repository at this point in the history
  • Loading branch information
fferrara committed Jun 8, 2018
1 parent ebb36c1 commit ff3332f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 5 deletions.
12 changes: 7 additions & 5 deletions fakeredis.py
Original file line number Diff line number Diff line change
Expand Up @@ -1627,8 +1627,8 @@ def _get_zelements_in_order(self, all_items, reverse=False):
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):
def zrangebyscore(self, name, min, max, 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 @@ -1638,11 +1638,13 @@ def zrangebyscore(self, name, min, max,
``withscores`` indicates to return the scores along with the values.
The return type is a list of (value, score) pairs
`score_cast_func`` a callable used to cast the score return value
"""
return self._zrangebyscore(name, min, max, start, num, withscores,
return self._zrangebyscore(name, min, max, start, num, withscores, score_cast_func,
reverse=False)

def _zrangebyscore(self, name, min, max, start, num, withscores, reverse):
def _zrangebyscore(self, name, min, max, start, num, withscores, score_cast_func, reverse):
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 @@ -1657,7 +1659,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 zrangebylex(self, name, min, max,
Expand Down
15 changes: 15 additions & 0 deletions test_fakeredis.py
Original file line number Diff line number Diff line change
Expand Up @@ -1879,6 +1879,21 @@ def test_zrangebyscore_withscores(self):
self.assertEqual(self.redis.zrangebyscore('foo', 1, 3, 0, 2, True),
[(b'one', 1), (b'two', 2)])

def test_zrangebyscore_cast_scores(self):
self.redis.zadd('foo', two=2)
self.redis.zadd('foo', two_a_also=2.2)

expected_without_cast_round = [(b'two', 2.0), (b'two_a_also', 2.2)]
expected_with_cast_round = [(b'two', 2.0), (b'two_a_also', 2.0)]
self.assertItemsEqual(
self.redis.zrangebyscore('foo', 2, 3, withscores=True),
expected_without_cast_round
)
self.assertItemsEqual(
self.redis.zrangebyscore('foo', 2, 3, withscores=True, score_cast_func=round),
expected_with_cast_round
)

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

0 comments on commit ff3332f

Please sign in to comment.