Add support for zrange arguments

fakeredis/_msgs.py

@@ -1,6 +1,8 @@
 INVALID_EXPIRE_MSG = "ERR invalid expire time in {}"
 WRONGTYPE_MSG = "WRONGTYPE Operation against a key holding the wrong kind of value"
 SYNTAX_ERROR_MSG = "ERR syntax error"
+SYNTAX_ERROR_LIMIT_ONLY_WITH_MSG = (
+    "ERR syntax error, LIMIT is only supported in combination with either BYSCORE or BYLEX")
 INVALID_INT_MSG = "ERR value is not an integer or out of range"
 INVALID_FLOAT_MSG = "ERR value is not a valid float"
 INVALID_OFFSET_MSG = "ERR offset is out of range"

fakeredis/commands_mixins/sortedset_mixin.py

@@ -156,17 +156,15 @@ 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 _zrange(self, key, start, stop, reverse, *args):
+    def _zrangebyscore(self, key, _min, _max, reverse, withscores, offset, count):
+        zset = key.value
+        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):
         zset = key.value
-        withscores = False
-        byscore = False
-        for arg in args:
-            if casematch(arg, b'withscores'):
-                withscores = True
-            elif casematch(arg, b'byscore'):
-                byscore = True
-            else:
-                raise SimpleError(msgs.SYNTAX_ERROR_MSG)
         if byscore:
             items = zset.irange_score(start.lower_bound, stop.upper_bound, reverse=reverse)
         else:
@@ -178,55 +176,74 @@ def _zrange(self, key, start, stop, reverse, *args):
         items = self._apply_withscores(items, withscores)
         return items
 
-    @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
-    def zrange(self, key, start, stop, *args):
-        return self._zrange(key, start, stop, False, *args)
-
-    @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
-    def zrevrange(self, key, start, stop, *args):
-        return self._zrange(key, start, stop, True, *args)
-
-    def _zrangebylex(self, key, _min, _max, reverse, *args):
-        if args:
-            if len(args) != 3 or not casematch(args[0], b'limit'):
-                raise SimpleError(msgs.SYNTAX_ERROR_MSG)
-            offset = Int.decode(args[1])
-            count = Int.decode(args[2])
-        else:
-            offset = 0
-            count = -1
+    def _zrangebylex(self, key, _min, _max, reverse, offset, count):
         zset = key.value
+        if reverse:
+            _min, _max = _max, _min
         items = zset.irange_lex(_min.value, _max.value,
                                 inclusive=(not _min.exclusive, not _max.exclusive),
                                 reverse=reverse)
         items = self._limit_items(items, offset, count)
         return items
 
+    def _zrange_args(self, key, start, stop, *args):
+        (bylex, byscore, rev, (offset, count), withscores), _ = extract_args(
+            args, ('bylex', 'byscore', 'rev', '++limit', 'withscores'))
+        if offset is not None and not bylex and not byscore:
+            raise SimpleError(msgs.SYNTAX_ERROR_LIMIT_ONLY_WITH_MSG)
+        if bylex and byscore:
+            raise SimpleError(msgs.SYNTAX_ERROR_MSG)
+
+        offset = offset or 0
+        count = -1 if count is None else count
+
+        if bylex:
+            res = self._zrangebylex(
+                key, StringTest.decode(start), StringTest.decode(stop), rev, offset, count)
+        elif byscore:
+            res = self._zrangebyscore(
+                key, ScoreTest.decode(start), ScoreTest.decode(stop), rev, withscores, offset, count)
+        else:
+            res = self._zrange(
+                key, ScoreTest.decode(start), ScoreTest.decode(stop), rev, withscores, byscore)
+        return res
+
+    @command((Key(ZSet), bytes, bytes), (bytes,))
+    def zrange(self, key, start, stop, *args):
+        return self._zrange_args(key, start, stop, *args)
+
+    @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
+    def zrevrange(self, key, start, stop, *args):
+        (withscores, byscore), _ = extract_args(args, ('withscores', 'byscore'))
+        return self._zrange(key, start, stop, True, withscores, byscore)
+
     @command((Key(ZSet), StringTest, StringTest), (bytes,))
     def zrangebylex(self, key, _min, _max, *args):
-        return self._zrangebylex(key, _min, _max, False, *args)
+        ((offset, count),), _ = extract_args(args, ('++limit',))
+        offset = offset or 0
+        count = -1 if count is None else count
+        return self._zrangebylex(key, _min, _max, False, offset, count)
 
     @command((Key(ZSet), StringTest, StringTest), (bytes,))
-    def zrevrangebylex(self, key, _max, _min, *args):
-        return self._zrangebylex(key, _min, _max, True, *args)
-
-    def _zrangebyscore(self, key, _min, _max, reverse, *args):
-        (withscores, (offset, count)), _ = extract_args(args, ('withscores', '++limit'))
+    def zrevrangebylex(self, key, _min, _max, *args):
+        ((offset, count),), _ = extract_args(args, ('++limit',))
         offset = offset or 0
         count = -1 if count is None else count
-        zset = key.value
-        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
+        return self._zrangebylex(key, _min, _max, True, offset, count)
 
     @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
     def zrangebyscore(self, key, _min, _max, *args):
-        return self._zrangebyscore(key, _min, _max, False, *args)
+        (withscores, (offset, count)), _ = extract_args(args, ('withscores', '++limit'))
+        offset = offset or 0
+        count = -1 if count is None else count
+        return self._zrangebyscore(key, _min, _max, False, withscores, offset, count)
 
     @command((Key(ZSet), ScoreTest, ScoreTest), (bytes,))
     def zrevrangebyscore(self, key, _max, _min, *args):
-        return self._zrangebyscore(key, _min, _max, True, *args)
+        (withscores, (offset, count)), _ = extract_args(args, ('withscores', '++limit'))
+        offset = offset or 0
+        count = -1 if count is None else count
+        return self._zrangebyscore(key, _min, _max, True, withscores, offset, count)
 
     @command((Key(ZSet), bytes))
     def zrank(self, key, member):

fakeredis/commands_mixins/string_mixin.py

@@ -89,7 +89,7 @@ def getdel(self, key):
         delete_keys(key)
         return res
 
-    @command((Key(bytes), Int, Int))
+    @command(name=['GETRANGE', 'SUBSTR'], fixed=(Key(bytes), Int, Int))
     def getrange(self, key, start, end):
         value = key.get(b'')
         start, end = fix_range_string(start, end, len(value))
@@ -216,11 +216,6 @@ def setrange(self, key, offset, value):
     def strlen(self, key):
         return len(key.get(b''))
 
-    # substr is a deprecated alias for getrange
-    @command((Key(bytes), Int, Int))
-    def substr(self, key, start, end):
-        return self.getrange(key, start, end)
-
     @command((Key(bytes),), (bytes,))
     def getex(self, key, *args):
         i, count_options, expire_time, diff = 0, 0, None, None

test/test_extract_args.py

@@ -54,6 +54,13 @@ def test_extract_args__multiple_numbers():
     assert limit == [324, 123]
     assert not keepttl
 
+    (xx, nx, limit, keepttl), _ = extract_args(
+        (b'nx', b'xx',),
+        ('nx', 'xx', '++limit', 'keepttl'))
+    assert xx
+    assert nx
+    assert not keepttl
+    assert limit == [None, None]
 
 def test_extract_args__extract_non_numbers():
     args = (b'by', b'dd', b'nx', b'limit', b'324', b'123', b'xx',)

test/test_mixins/test_sortedset_commands.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import math
 from collections import OrderedDict
+from typing import Tuple, List, Optional
 
-import math
 import pytest
 import redis
 import redis.client
 from packaging.version import Version
-from typing import Tuple, List, Optional
+
+from test import testtools
 
 REDIS_VERSION = Version(redis.__version__)
 
@@ -53,6 +55,62 @@ def test_zrange_same_score(r):
     assert r.zrange('foo', 2, 3) == [b'two_c', b'two_d']
 
 
+def test_zrange_with_bylex_and_byscore(r: redis.Redis):
+    r.zadd('foo', {'one_a': 0})
+    r.zadd('foo', {'two_a': 0})
+    r.zadd('foo', {'two_b': 0})
+    r.zadd('foo', {'three_a': 0})
+    with pytest.raises(redis.ResponseError):
+        testtools.raw_command(r, 'zrange', 'foo', '(t', '+', 'bylex', 'byscore')
+
+
+def test_zrange_with_rev_and_bylex(r: redis.Redis):
+    r.zadd('foo', {'one_a': 0})
+    r.zadd('foo', {'two_a': 0})
+    r.zadd('foo', {'two_b': 0})
+    r.zadd('foo', {'three_a': 0})
+    assert r.zrange('foo', b'+', b'(t', desc=True, bylex=True) == [b'two_b', b'two_a', b'three_a']
+    assert (
+            r.zrange('foo', b'[two_b', b'(t', desc=True, bylex=True)
+            == [b'two_b', b'two_a', b'three_a']
+    )
+    assert r.zrange('foo', b'(two_b', b'(t', desc=True, bylex=True) == [b'two_a', b'three_a']
+    assert (
+            r.zrange('foo', b'[two_b', b'[three_a', desc=True, bylex=True)
+            == [b'two_b', b'two_a', b'three_a']
+    )
+    assert r.zrange('foo', b'[two_b', b'(three_a', desc=True, bylex=True) == [b'two_b', b'two_a']
+    assert r.zrange('foo', b'(two_b', b'-', desc=True, bylex=True) == [b'two_a', b'three_a', b'one_a']
+    assert r.zrange('foo', b'(two_b', b'[two_b', bylex=True) == []
+    # reversed max + and min - boundaries
+    # these will be always empty, but allowed by redis
+    assert r.zrange('foo', b'-', b'+', desc=True, bylex=True) == []
+    assert r.zrange('foo', b'[three_a', b'+', desc=True, bylex=True) == []
+    assert r.zrange('foo', b'-', b'[o', desc=True, bylex=True) == []
+
+
+def test_zrange_with_bylex(r):
+    r.zadd('foo', {'one_a': 0})
+    r.zadd('foo', {'two_a': 0})
+    r.zadd('foo', {'two_b': 0})
+    r.zadd('foo', {'three_a': 0})
+    assert r.zrange('foo', b'(t', b'+', bylex=True) == [b'three_a', b'two_a', b'two_b']
+    assert r.zrange('foo', b'(t', b'[two_b', bylex=True) == [b'three_a', b'two_a', b'two_b']
+    assert r.zrange('foo', b'(t', b'(two_b', bylex=True) == [b'three_a', b'two_a']
+    assert (
+            r.zrange('foo', b'[three_a', b'[two_b', bylex=True)
+            == [b'three_a', b'two_a', b'two_b']
+    )
+    assert r.zrange('foo', b'(three_a', b'[two_b', bylex=True) == [b'two_a', b'two_b']
+    assert r.zrange('foo', b'-', b'(two_b', bylex=True) == [b'one_a', b'three_a', b'two_a']
+    assert r.zrange('foo', b'[two_b', b'(two_b', bylex=True) == []
+    # reversed max + and min - boundaries
+    # these will be always empty, but allowed by redis
+    assert r.zrange('foo', b'+', b'-', bylex=True) == []
+    assert r.zrange('foo', b'+', b'[three_a', bylex=True) == []
+    assert r.zrange('foo', b'[o', b'-', bylex=True) == []
+
+
 def test_zrange_with_byscore(r):
     r.zadd('foo', {'zero': 0})
     r.zadd('foo', {'two': 2})