From f284900d50d39aedc3db8f733ffd46cda0f444c5 Mon Sep 17 00:00:00 2001 From: Flavio Ferrara Date: Mon, 11 Jun 2018 13:41:41 +0100 Subject: [PATCH] Supports argument score_cast_func in zrange(), zrevrange() --- fakeredis.py | 16 ++++++++++------ test_fakeredis.py | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/fakeredis.py b/fakeredis.py index 3203767..97ba248 100644 --- a/fakeredis.py +++ b/fakeredis.py @@ -1593,7 +1593,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. @@ -1604,6 +1604,8 @@ def zrange(self, name, start, end, desc=False, withscores=False): ``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 """ if end == -1: end = None @@ -1619,7 +1621,7 @@ 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( @@ -1639,7 +1641,7 @@ def zrangebyscore(self, name, min, max, start=None, num=None, ``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 + ``score_cast_func`` a callable used to cast the score return value """ return self._zrangebyscore(name, min, max, start, num, withscores, score_cast_func, reverse=False) @@ -1797,7 +1799,7 @@ def zlexcount(self, name, min, max): found += 1 return found - 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. @@ -1806,8 +1808,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 + + ``score_cast_func`` a callable used to cast the score return value """ - 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, score_cast_func=float): @@ -1821,7 +1825,7 @@ def zrevrangebyscore(self, name, max, min, start=None, num=None, ``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 + ``score_cast_func`` a callable used to cast the score return value """ return self._zrangebyscore(name, min, max, start, num, withscores, score_cast_func, reverse=True) diff --git a/test_fakeredis.py b/test_fakeredis.py index 9b5ac62..049c5f2 100644 --- a/test_fakeredis.py +++ b/test_fakeredis.py @@ -1716,6 +1716,17 @@ def test_zrange_wrong_type(self): with self.assertRaises(redis.ResponseError): self.redis.zrange('foo', 0, -1) + def test_zrange_score_cast(self): + self.redis.zadd('foo', one=1.2) + self.redis.zadd('foo', two=2.2) + + expected_without_cast_round = [(b'one', 1.2), (b'two', 2.2)] + expected_with_cast_round = [(b'one', 1.0), (b'two', 2.0)] + self.assertEqual(self.redis.zrange('foo', 0, 2, withscores=True), + expected_without_cast_round) + self.assertEqual(self.redis.zrange('foo', 0, 2, withscores=True, score_cast_func=round), + expected_with_cast_round) + def test_zrank(self): self.redis.zadd('foo', one=1) self.redis.zadd('foo', two=2) @@ -1811,6 +1822,17 @@ def test_zrevrange_wrong_type(self): with self.assertRaises(redis.ResponseError): self.redis.zrevrange('foo', 0, 2) + def test_zrevrange_score_cast(self): + self.redis.zadd('foo', one=1.2) + self.redis.zadd('foo', two=2.2) + + expected_without_cast_round = [(b'two', 2.2), (b'one', 1.2)] + expected_with_cast_round = [(b'two', 2.0), (b'one', 1.0)] + self.assertEqual(self.redis.zrevrange('foo', 0, 2, withscores=True), + expected_without_cast_round) + self.assertEqual(self.redis.zrevrange('foo', 0, 2, withscores=True, score_cast_func=round), + expected_with_cast_round) + def test_zrangebyscore(self): self.redis.zadd('foo', zero=0) self.redis.zadd('foo', two=2) @@ -1945,6 +1967,24 @@ def test_zrevrangebyscore_wrong_type(self): with self.assertRaises(redis.ResponseError): self.redis.zrevrangebyscore('foo', '(3', '(1') + def test_zrevrangebyscore_cast_scores(self): + self.redis.zadd('foo', two=2) + self.redis.zadd('foo', two_a_also=2.2) + + def round_str(x): + return round(float(x)) + + expected_without_cast_round = [(b'two_a_also', 2.2), (b'two', 2.0)] + expected_with_cast_round = [(b'two_a_also', 2.0), (b'two', 2.0)] + self.assertEqual( + self.redis.zrevrangebyscore('foo', 3, 2, withscores=True), + expected_without_cast_round + ) + self.assertEqual( + self.redis.zrevrangebyscore('foo', 3, 2, withscores=True, score_cast_func=round_str), + expected_with_cast_round + ) + def test_zrangebylex(self): self.redis.zadd('foo', one_a=0) self.redis.zadd('foo', two_a=0)