Supports argument score_cast_func in zrange(), zrevrange()

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)

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)