From ff3332f5445ffac46b9c4a5c46671243581be316 Mon Sep 17 00:00:00 2001
From: Flavio Ferrara <femferrara@gmail.com>
Date: Fri, 8 Jun 2018 16:03:06 +0100
Subject: [PATCH] Supports argument score_cast_func in zrangebyscore()

---
 fakeredis.py      | 12 +++++++-----
 test_fakeredis.py | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/fakeredis.py b/fakeredis.py
index 03c9393..d7552cb 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,
diff --git a/test_fakeredis.py b/test_fakeredis.py
index bd46528..5aeb9d5 100644
--- a/test_fakeredis.py
+++ b/test_fakeredis.py
@@ -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)