feat[sorted-set]:implement ZRANGESTORE

Fix #193

docs/about/changelog.md

@@ -13,7 +13,7 @@ description: Change log of all fakeredis releases
     `XACK` #157, `XPENDING` #170
 - Implemented sorted set commands: 
   - `ZRANDMEMBER` #192, `ZDIFF` #187, `ZINTER` #189, `ZUNION` #194, `ZDIFFSTORE` #188,
-    `ZINTERCARD` #190
+    `ZINTERCARD` #190, `ZRANGESTORE` #193
 
 ### 🧰 Maintenance
 

docs/redis-commands/Redis.md

@@ -827,6 +827,10 @@ Returns members in a sorted set within a lexicographical range.
 
 Returns members in a sorted set within a range of scores.
 
+### [ZRANGESTORE](https://redis.io/commands/zrangestore/)
+
+Stores a range of members from sorted set in a key.
+
 ### [ZRANK](https://redis.io/commands/zrank/)
 
 Returns the index of a member in a sorted set ordered by ascending scores.
@@ -891,10 +895,6 @@ Removes and returns a member by score from one or more sorted sets. Blocks until
 
 Returns the highest- or lowest-scoring members from one or more sorted sets after removing them. Deletes the sorted set if the last member was popped.
 
-#### [ZRANGESTORE](https://redis.io/commands/zrangestore/) <small>(not implemented)</small>
-
-Stores a range of members from sorted set in a key.
-
 
 ## generic commands
 

fakeredis/_commands.py

@@ -5,7 +5,7 @@
 import functools
 import math
 import re
-from typing import Tuple
+from typing import Tuple, Union
 
 from . import _msgs as msgs
 from ._helpers import null_terminate, SimpleError, SimpleString
@@ -251,7 +251,7 @@ def __eq__(self, other):
 class ScoreTest:
     """Argument converter for sorted set score endpoints."""
 
-    def __init__(self, value, exclusive=False, bytes_val=None):
+    def __init__(self, value: float, exclusive: bool = False, bytes_val: bytes = None):
         self.value = value
         self.exclusive = exclusive
         self.bytes_val = bytes_val
@@ -289,12 +289,12 @@ def upper_bound(self):
 class StringTest:
     """Argument converter for sorted set LEX endpoints."""
 
-    def __init__(self, value, exclusive):
+    def __init__(self, value: Union[bytes, BeforeAny, AfterAny], exclusive: bool):
         self.value = value
         self.exclusive = exclusive
 
     @classmethod
-    def decode(cls, value):
+    def decode(cls, value: bytes) -> 'StringTest':
         if value == b'-':
             return cls(BeforeAny(), True)
         elif value == b'+':
@@ -306,6 +306,13 @@ def decode(cls, value):
         else:
             raise SimpleError(msgs.INVALID_MIN_MAX_STR_MSG)
 
+    def to_scoretest(self, zset: ZSet):
+        if isinstance(self.value, BeforeAny):
+            return ScoreTest(float('-inf'), False)
+        if isinstance(self.value, AfterAny):
+            return ScoreTest(float('inf'), False)
+        return ScoreTest(zset.get(self.value, None), self.exclusive)
+
 
 class Signature:
     def __init__(self, name, func_name, fixed, repeat=(), args=(), flags=""):

fakeredis/_zset.py

@@ -1,3 +1,5 @@
+from typing import Any, Tuple
+
 import sortedcontainers
 
 
@@ -77,7 +79,7 @@ def irange_lex(self, start, stop, inclusive=(True, True), reverse=False):
                                   inclusive=inclusive, reverse=reverse)
         return (item[1] for item in it)
 
-    def irange_score(self, start, stop, reverse=False):
+    def irange_score(self, start: Tuple[Any, bytes], stop: Tuple[Any, bytes], reverse:bool=False):
         return self._byscore.irange(start, stop, reverse=reverse)
 
     def rank(self, member):

fakeredis/commands_mixins/sortedset_mixin.py

@@ -5,7 +5,7 @@
 import math
 import random
 import sys
-from typing import Union, Optional
+from typing import Union, Optional, List
 
 from fakeredis import _msgs as msgs
 from fakeredis._command_args_parsing import extract_args
@@ -78,7 +78,7 @@ def _limit_items(items, offset, count):
             out.append(item)
         return out
 
-    def _apply_withscores(self, items, withscores):
+    def _apply_withscores(self, items, withscores: bool) -> List[bytes]:
         if withscores:
             out = []
             for item in items:
@@ -168,14 +168,16 @@ def zincrby(self, key, increment, member):
     def zlexcount(self, key, _min, _max):
         return key.value.zlexcount(_min.value, _min.exclusive, _max.value, _max.exclusive)
 
-    def _zrangebyscore(self, key, _min, _max, reverse, withscores, offset, count):
+    def _zrangebyscore(self, key, _min, _max, reverse, withscores, offset, count) -> List[bytes]:
         zset = key.value
+        if reverse:
+            _min, _max = _max, _min
         items = list(zset.irange_score(_min.lower_bound, _max.upper_bound, reverse=reverse))
         items = self._limit_items(items, offset, count)
         items = self._apply_withscores(items, withscores)
         return items
 
-    def _zrange(self, key, start, stop, reverse, withscores, byscore):
+    def _zrange(self, key, start, stop, reverse, withscores, byscore) -> List[bytes]:
         zset = key.value
         if byscore:
             items = zset.irange_score(start.lower_bound, stop.upper_bound, reverse=reverse)
@@ -188,7 +190,7 @@ def _zrange(self, key, start, stop, reverse, withscores, byscore):
         items = self._apply_withscores(items, withscores)
         return items
 
-    def _zrangebylex(self, key, _min, _max, reverse, offset, count):
+    def _zrangebylex(self, key, _min, _max, reverse, offset, count) -> List[bytes]:
         zset = key.value
         if reverse:
             _min, _max = _max, _min
@@ -224,6 +226,15 @@ def _zrange_args(self, key, start, stop, *args):
     def zrange(self, key, start, stop, *args):
         return self._zrange_args(key, start, stop, *args)
 
+    @command((Key(ZSet), Key(ZSet), bytes, bytes), (bytes,))
+    def zrangestore(self, dest: CommandItem, src, start, stop, *args):
+        results_list = self._zrange_args(src, start, stop, *args)
+        res = ZSet()
+        for item in results_list:
+            res.add(item, src.value.get(item))
+        dest.update(res)
+        return len(res)
+
     @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
     def zrevrange(self, key, start, stop, *args):
         (withscores, byscore), _ = extract_args(args, ('withscores', 'byscore'))
@@ -251,7 +262,7 @@ def zrangebyscore(self, key, _min, _max, *args):
         return self._zrangebyscore(key, _min, _max, False, withscores, offset, count)
 
     @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
-    def zrevrangebyscore(self, key, _max, _min, *args):
+    def zrevrangebyscore(self, key, _min, _max, *args):
         (withscores, (offset, count)), _ = extract_args(args, ('withscores', '++limit'))
         offset = offset or 0
         count = -1 if count is None else count
@@ -283,8 +294,8 @@ def zrem(self, key, *members):
 
     @command((Key(ZSet), StringTest, StringTest))
     def zremrangebylex(self, key, _min, _max):
-        items = key.value.irange_lex(_min.value, _max.value,
-                                     inclusive=(not _min.exclusive, not _max.exclusive))
+        items = key.value.irange_lex(
+            _min.value, _max.value, inclusive=(not _min.exclusive, not _max.exclusive))
         return self.zrem(key, *items)
 
     @command((Key(ZSet), ScoreTest, ScoreTest))

test/test_mixins/test_sortedset_commands.py

@@ -553,17 +553,10 @@ def test_zrevrangebyscore_cast_scores(r: redis.Redis):
     r.zadd('foo', {'two': 2})
     r.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)]
-    assert (
-            r.zrevrangebyscore('foo', 3, 2, withscores=True)
-            == expected_without_cast_round
-    )
-    assert (
-            r.zrevrangebyscore('foo', 3, 2, withscores=True,
-                               score_cast_func=round_str)
-            == expected_with_cast_round
-    )
+    assert r.zrevrangebyscore('foo', 3, 2, withscores=True) == [(b'two_a_also', 2.2), (b'two', 2.0)]
+
+    assert r.zrevrangebyscore(
+        'foo', 3, 2, withscores=True, score_cast_func=round_str) == [(b'two_a_also', 2.0), (b'two', 2.0)]
 
 
 def test_zrangebylex(r: redis.Redis):
@@ -1155,3 +1148,23 @@ def test_zintercard(r: redis.Redis):
     r.zadd("c", {"a1": 6, "a3": 5, "a4": 4})
     assert r.zintercard(3, ["a", "b", "c"]) == 2
     assert r.zintercard(3, ["a", "b", "c"], limit=1) == 1
+
+
+def test_zrangestore(r: redis.Redis):
+    r.zadd("a", {"a1": 1, "a2": 2, "a3": 3})
+    assert r.zrangestore("b", "a", 0, 1)
+    assert r.zrange("b", 0, -1) == [b"a1", b"a2"]
+    assert r.zrangestore("b", "a", 1, 2)
+    assert r.zrange("b", 0, -1) == [b"a2", b"a3"]
+    assert r.zrange("b", 0, -1, withscores=True) == [(b"a2", 2), (b"a3", 3)]
+    # reversed order
+    assert r.zrangestore("b", "a", 1, 2, desc=True)
+    assert r.zrange("b", 0, -1) == [b"a1", b"a2"]
+    # by score
+    assert r.zrangestore("b", "a", 2, 1, byscore=True, offset=0, num=1, desc=True)
+    assert r.zrange("b", 0, -1) == [b"a2"]
+    # by lex
+    # TODO: fix
+    # assert r.zrange("a", "[a2", "(a3", bylex=True, offset=0, num=1) == [b"a2"]
+    # assert r.zrangestore("b", "a", "[a2", "(a3", bylex=True, offset=0, num=1)
+    # assert r.zrange("b", 0, -1) == [b"a2"]