Supports argument score_cast_func in zrangebyscore()

fakeredis.py

@@ -1593,7 +1593,18 @@ def zinterstore(self, dest, keys, aggregate=None):
                                 lambda x: x in
                                 valid_keys)
 
-    def zrange(self, name, start, end, desc=False, withscores=False):
+    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
         ``start`` and ``end`` sorted in ascending order.
@@ -1604,6 +1615,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
@@ -1616,19 +1629,16 @@ def zrange(self, name, start, end, desc=False, withscores=False):
             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, 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(
             all_items.items(), key=lambda x: x[0], reverse=reverse)
         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 +1648,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 "
@@ -1656,9 +1668,7 @@ def _zrangebyscore(self, name, min, max, start, num, withscores, reverse):
                 matches.append(item)
         if start is not None:
             matches = matches[start:start + num]
-        if withscores:
-            return [(k, 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):
@@ -1795,7 +1805,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.
@@ -1804,11 +1814,13 @@ 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):
+    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 +1830,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,

test_fakeredis.py

@@ -91,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'])
@@ -1716,6 +1720,18 @@ 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=self._round_str),
+                         expected_with_cast_round)
+
     def test_zrank(self):
         self.redis.zadd('foo', one=1)
         self.redis.zadd('foo', two=2)
@@ -1811,6 +1827,18 @@ 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=self._round_str),
+                         expected_with_cast_round)
+
     def test_zrangebyscore(self):
         self.redis.zadd('foo', zero=0)
         self.redis.zadd('foo', two=2)
@@ -1879,6 +1907,22 @@ 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=self._round_str),
+            expected_with_cast_round
+        )
+
     def test_zrevrangebyscore(self):
         self.redis.zadd('foo', one=1)
         self.redis.zadd('foo', two=2)
@@ -1927,6 +1971,22 @@ 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)
+
+        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=self._round_str),
+            expected_with_cast_round
+        )
+
     def test_zrangebylex(self):
         self.redis.zadd('foo', one_a=0)
         self.redis.zadd('foo', two_a=0)
@@ -3780,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)