From fdff224f2df62c3885349296f8a833b9c67fbdf9 Mon Sep 17 00:00:00 2001 From: Flavio Ferrara Date: Fri, 8 Jun 2018 16:03:06 +0100 Subject: [PATCH 1/4] Supports argument score_cast_func in zrangebyscore() --- fakeredis.py | 20 ++++++++++++-------- test_fakeredis.py | 18 ++++++++++++++++++ 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/fakeredis.py b/fakeredis.py index 03c9393..3203767 100644 --- a/fakeredis.py +++ b/fakeredis.py @@ -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``. @@ -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 " @@ -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, @@ -1807,8 +1809,8 @@ def zrevrange(self, name, start, num, withscores=False): """ return self.zrange(name, start, num, True, withscores) - def zrevrangebyscore(self, name, max, min, - start=None, num=None, withscores=False): + def zrevrangebyscore(self, name, max, min, 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. @@ -1818,8 +1820,10 @@ def zrevrangebyscore(self, name, max, min, ``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=True) def zrevrangebylex(self, name, max, min, diff --git a/test_fakeredis.py b/test_fakeredis.py index bd46528..9b5ac62 100644 --- a/test_fakeredis.py +++ b/test_fakeredis.py @@ -1879,6 +1879,24 @@ 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) + + def round_str(x): + return round(float(x)) + + 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_str), + expected_with_cast_round + ) + def test_zrevrangebyscore(self): self.redis.zadd('foo', one=1) self.redis.zadd('foo', two=2) From b0eea8dbeac2aa03901d1c73a45f32fbb8f0cab4 Mon Sep 17 00:00:00 2001 From: Flavio Ferrara Date: Mon, 11 Jun 2018 13:41:41 +0100 Subject: [PATCH 2/4] Supports argument score_cast_func in zrange(), zrevrange() --- fakeredis.py | 16 ++++++++++------ test_fakeredis.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 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..bd108bc 100644 --- a/test_fakeredis.py +++ b/test_fakeredis.py @@ -1716,6 +1716,20 @@ 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) + + def round_str(x): + return round(float(x)) + + 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_str), + expected_with_cast_round) + def test_zrank(self): self.redis.zadd('foo', one=1) self.redis.zadd('foo', two=2) @@ -1811,6 +1825,21 @@ 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) + + def round_str(x): + return round(float(x)) + + 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_str), + expected_with_cast_round) + def test_zrangebyscore(self): self.redis.zadd('foo', zero=0) self.redis.zadd('foo', two=2) @@ -1945,6 +1974,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) From 777e25e4d49d2682e677434cd360c70bf97f01ce Mon Sep 17 00:00:00 2001 From: Flavio Ferrara Date: Tue, 12 Jun 2018 16:11:12 +0100 Subject: [PATCH 3/4] Scores are converted to bytes in Fakeredis. Tests fixed --- fakeredis.py | 6 ++++-- test_fakeredis.py | 19 +++++++------------ 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/fakeredis.py b/fakeredis.py index 97ba248..e6fe113 100644 --- a/fakeredis.py +++ b/fakeredis.py @@ -1621,7 +1621,8 @@ def zrange(self, name, start, end, desc=False, withscores=False, score_cast_func if not withscores: return items else: - return [(k, score_cast_func(all_items[k])) for k in items] + return [(k, score_cast_func(to_bytes(all_items[k]))) + for k in items] def _get_zelements_in_order(self, all_items, reverse=False): by_keyname = sorted( @@ -1661,7 +1662,8 @@ def _zrangebyscore(self, name, min, max, start, num, withscores, score_cast_func if start is not None: matches = matches[start:start + num] if withscores: - return [(k, score_cast_func(all_items[k])) for k in matches] + return [(k, score_cast_func(to_bytes(all_items[k]))) + for k in matches] return matches def zrangebylex(self, name, min, max, diff --git a/test_fakeredis.py b/test_fakeredis.py index bd108bc..929db52 100644 --- a/test_fakeredis.py +++ b/test_fakeredis.py @@ -74,6 +74,13 @@ def key_val_dict(size=100): for i in range(size)]) +def round_str(x): + if not (isinstance(x, str) or hasattr(x, 'decode')): + raise AssertionError('Cast argument should be str or bytes.') + + return round(float(x)) + + class TestFakeStrictRedis(unittest.TestCase): decode_responses = False @@ -1720,9 +1727,6 @@ def test_zrange_score_cast(self): self.redis.zadd('foo', one=1.2) self.redis.zadd('foo', two=2.2) - def round_str(x): - return round(float(x)) - 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), @@ -1829,9 +1833,6 @@ def test_zrevrange_score_cast(self): self.redis.zadd('foo', one=1.2) self.redis.zadd('foo', two=2.2) - def round_str(x): - return round(float(x)) - 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), @@ -1912,9 +1913,6 @@ def test_zrangebyscore_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', 2.0), (b'two_a_also', 2.2)] expected_with_cast_round = [(b'two', 2.0), (b'two_a_also', 2.0)] self.assertItemsEqual( @@ -1978,9 +1976,6 @@ 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( From aa6ff5426a3f63bed896c0435a8a06534d41462b Mon Sep 17 00:00:00 2001 From: Bruce Merry Date: Fri, 15 Jun 2018 10:26:45 +0200 Subject: [PATCH 4/4] Handle decode_responses correctly for score_cast_func When decode_responses is true, the score_cast_func is passed text, not bytes. To handle this, the score_cast_func handling code was unified in one helper function which deals with the different cases. It also has a cut-through path for the default score_cast_func to avoid a lot of unnecessarily float -> bytes -> text -> float casting. The unit test is made stricted to check the type exactly. --- fakeredis.py | 22 +++++++++++++--------- test_fakeredis.py | 26 +++++++++++++++----------- 2 files changed, 28 insertions(+), 20 deletions(-) diff --git a/fakeredis.py b/fakeredis.py index e6fe113..470597f 100644 --- a/fakeredis.py +++ b/fakeredis.py @@ -1593,6 +1593,17 @@ def zinterstore(self, dest, keys, aggregate=None): lambda x: x in valid_keys) + def _apply_score_cast_func(self, items, all_items, withscores, score_cast_func): + if not withscores: + return items + elif score_cast_func is float: + # Fast path for common case + return [(k, all_items[k]) for k in items] + elif self._decode_responses: + return [(k, score_cast_func(_decode(to_bytes(all_items[k])))) for k in items] + else: + return [(k, score_cast_func(to_bytes(all_items[k]))) for k in items] + def zrange(self, name, start, end, desc=False, withscores=False, score_cast_func=float): """ Return a range of values from sorted set ``name`` between @@ -1618,11 +1629,7 @@ def zrange(self, name, start, end, desc=False, withscores=False, score_cast_func reverse = False in_order = self._get_zelements_in_order(all_items, reverse) items = in_order[start:end] - if not withscores: - return items - else: - return [(k, score_cast_func(to_bytes(all_items[k]))) - for k in items] + return self._apply_score_cast_func(items, all_items, withscores, score_cast_func) def _get_zelements_in_order(self, all_items, reverse=False): by_keyname = sorted( @@ -1661,10 +1668,7 @@ def _zrangebyscore(self, name, min, max, start, num, withscores, score_cast_func matches.append(item) if start is not None: matches = matches[start:start + num] - if withscores: - return [(k, score_cast_func(to_bytes(all_items[k]))) - for k in matches] - return matches + return self._apply_score_cast_func(matches, all_items, withscores, score_cast_func) def zrangebylex(self, name, min, max, start=None, num=None): diff --git a/test_fakeredis.py b/test_fakeredis.py index 929db52..f08527b 100644 --- a/test_fakeredis.py +++ b/test_fakeredis.py @@ -74,13 +74,6 @@ def key_val_dict(size=100): for i in range(size)]) -def round_str(x): - if not (isinstance(x, str) or hasattr(x, 'decode')): - raise AssertionError('Cast argument should be str or bytes.') - - return round(float(x)) - - class TestFakeStrictRedis(unittest.TestCase): decode_responses = False @@ -98,6 +91,10 @@ def assertItemsEqual(self, a, b): def create_redis(self, db=0): return fakeredis.FakeStrictRedis(db=db) + def _round_str(self, x): + self.assertIsInstance(x, bytes) + return round(float(x)) + def test_flushdb(self): self.redis.set('foo', 'bar') self.assertEqual(self.redis.keys(), [b'foo']) @@ -1731,7 +1728,8 @@ def test_zrange_score_cast(self): 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_str), + self.assertEqual(self.redis.zrange('foo', 0, 2, withscores=True, + score_cast_func=self._round_str), expected_with_cast_round) def test_zrank(self): @@ -1838,7 +1836,7 @@ def test_zrevrange_score_cast(self): 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_str), + score_cast_func=self._round_str), expected_with_cast_round) def test_zrangebyscore(self): @@ -1920,7 +1918,8 @@ def test_zrangebyscore_cast_scores(self): expected_without_cast_round ) self.assertItemsEqual( - self.redis.zrangebyscore('foo', 2, 3, withscores=True, score_cast_func=round_str), + self.redis.zrangebyscore('foo', 2, 3, withscores=True, + score_cast_func=self._round_str), expected_with_cast_round ) @@ -1983,7 +1982,8 @@ def test_zrevrangebyscore_cast_scores(self): expected_without_cast_round ) self.assertEqual( - self.redis.zrevrangebyscore('foo', 3, 2, withscores=True, score_cast_func=round_str), + self.redis.zrevrangebyscore('foo', 3, 2, withscores=True, + score_cast_func=self._round_str), expected_with_cast_round ) @@ -3840,6 +3840,10 @@ def test_lock(self): class DecodeMixin(object): decode_responses = True + def _round_str(self, x): + self.assertIsInstance(x, fakeredis.text_type) + return round(float(x)) + def assertEqual(self, a, b, msg=None): super(DecodeMixin, self).assertEqual(a, fakeredis._decode(b), msg)