add support to `ZRANGE` and `ZRANGESTORE` parameters (#1603)

2 files changed, 181 insertions, 122 deletions

diff --git a/redis/commands.py b/redis/commands.py
index baf5239..ee883bc 100644
--- a/redis/commands.py
+++ b/redis/commands.py
@@ -1073,7 +1073,7 @@ class Commands:
         return self.execute_command("HRANDFIELD", key, *params)
 
     def randomkey(self):
-        "Returns the name of a random key"
+        """Returns the name of a random key"""
         return self.execute_command('RANDOMKEY')
 
     def rename(self, src, dst):
@@ -1083,7 +1083,7 @@ class Commands:
         return self.execute_command('RENAME', src, dst)
 
     def renamenx(self, src, dst):
-        "Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist"
+        """Rename key ``src`` to ``dst`` if ``dst`` doesn't already exist"""
         return self.execute_command('RENAMENX', src, dst)
 
     def restore(self, name, ttl, value, replace=False, absttl=False,
@@ -1545,32 +1545,25 @@ class Commands:
 
         pieces = [name]
         if by is not None:
-            pieces.append(b'BY')
-            pieces.append(by)
+            pieces.extend([b'BY', by])
         if start is not None and num is not None:
-            pieces.append(b'LIMIT')
-            pieces.append(start)
-            pieces.append(num)
+            pieces.extend([b'LIMIT', start, num])
         if get is not None:
             # If get is a string assume we want to get a single value.
             # Otherwise assume it's an interable and we want to get multiple
             # values. We can't just iterate blindly because strings are
             # iterable.
             if isinstance(get, (bytes, str)):
-                pieces.append(b'GET')
-                pieces.append(get)
+                pieces.extend([b'GET', get])
             else:
                 for g in get:
-                    pieces.append(b'GET')
-                    pieces.append(g)
+                    pieces.extend([b'GET', g])
         if desc:
             pieces.append(b'DESC')
         if alpha:
             pieces.append(b'ALPHA')
         if store is not None:
-            pieces.append(b'STORE')
-            pieces.append(store)
-
+            pieces.extend([b'STORE', store])
         if groups:
             if not get or isinstance(get, (bytes, str)) or len(get) < 2:
                 raise DataError('when using "groups" the "get" argument '
@@ -1729,15 +1722,15 @@ class Commands:
 
     # SET COMMANDS
     def sadd(self, name, *values):
-        "Add ``value(s)`` to set ``name``"
+        """Add ``value(s)`` to set ``name``"""
         return self.execute_command('SADD', name, *values)
 
     def scard(self, name):
-        "Return the number of elements in set ``name``"
+        """Return the number of elements in set ``name``"""
         return self.execute_command('SCARD', name)
 
     def sdiff(self, keys, *args):
-        "Return the difference of sets specified by ``keys``"
+        """Return the difference of sets specified by ``keys``"""
         args = list_or_args(keys, args)
         return self.execute_command('SDIFF', *args)
 
@@ -1750,7 +1743,7 @@ class Commands:
         return self.execute_command('SDIFFSTORE', dest, *args)
 
     def sinter(self, keys, *args):
-        "Return the intersection of sets specified by ``keys``"
+        """Return the intersection of sets specified by ``keys``"""
         args = list_or_args(keys, args)
         return self.execute_command('SINTER', *args)
 
@@ -1763,15 +1756,17 @@ class Commands:
         return self.execute_command('SINTERSTORE', dest, *args)
 
     def sismember(self, name, value):
-        "Return a boolean indicating if ``value`` is a member of set ``name``"
+        """
+        Return a boolean indicating if ``value`` is a member of set ``name``
+        """
         return self.execute_command('SISMEMBER', name, value)
 
     def smembers(self, name):
-        "Return all members of the set ``name``"
+        """Return all members of the set ``name``"""
         return self.execute_command('SMEMBERS', name)
 
     def smove(self, src, dst, value):
-        "Move ``value`` from set ``src`` to set ``dst`` atomically"
+        """Move ``value`` from set ``src`` to set ``dst`` atomically"""
         return self.execute_command('SMOVE', src, dst, value)
 
     def spop(self, name, count=None):
@@ -1850,8 +1845,7 @@ class Commands:
                 pieces.append(b'~')
             pieces.append(minid)
         if limit is not None:
-            pieces.append(b"LIMIT")
-            pieces.append(limit)
+            pieces.extend([b'LIMIT', limit])
         if nomkstream:
             pieces.append(b'NOMKSTREAM')
         pieces.append(id)
@@ -2440,41 +2434,113 @@ class Commands:
         keys.append(timeout)
         return self.execute_command('BZPOPMIN', *keys)
 
+    def _zrange(self, command, dest, name, start, end, desc=False,
+                byscore=False, bylex=False, withscores=False,
+                score_cast_func=float, offset=None, num=None):
+        if byscore and bylex:
+            raise DataError("``byscore`` and ``bylex`` can not be "
+                            "specified together.")
+        if (offset is not None and num is None) or \
+                (num is not None and offset is None):
+            raise DataError("``offset`` and ``num`` must both be specified.")
+        if bylex and withscores:
+            raise DataError("``withscores`` not supported in combination "
+                            "with ``bylex``.")
+        pieces = [command]
+        if dest:
+            pieces.append(dest)
+        pieces.extend([name, start, end])
+        if byscore:
+            pieces.append('BYSCORE')
+        if bylex:
+            pieces.append('BYLEX')
+        if desc:
+            pieces.append('REV')
+        if offset is not None and num is not None:
+            pieces.extend(['LIMIT', offset, num])
+        if withscores:
+            pieces.append('WITHSCORES')
+        options = {
+            'withscores': withscores,
+            'score_cast_func': score_cast_func
+        }
+        return self.execute_command(*pieces, **options)
+
     def zrange(self, name, start, end, desc=False, withscores=False,
-               score_cast_func=float):
+               score_cast_func=float, byscore=False, bylex=False,
+               offset=None, num=None):
         """
         Return a range of values from sorted set ``name`` between
         ``start`` and ``end`` sorted in ascending order.
 
         ``start`` and ``end`` can be negative, indicating the end of the range.
 
-        ``desc`` a boolean indicating whether to sort the results descendingly
+        ``desc`` a boolean indicating whether to sort the results in reversed
+        order.
 
         ``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.
+
+        ``byscore`` when set to True, returns the range of elements from the
+        sorted set having scores equal or between ``start`` and ``end``.
+
+        ``bylex`` when set to True, returns the range of elements from the
+        sorted set between the ``start`` and ``end`` lexicographical closed
+        range intervals.
+        Valid ``start`` and ``end`` must start with ( or [, in order to specify
+        whether the range interval is exclusive or inclusive, respectively.
+
+        ``offset`` and ``num`` are specified, then return a slice of the range.
+        Can't be provided when using ``bylex``.
+        """
+        return self._zrange('ZRANGE', None, name, start, end, desc, byscore,
+                            bylex, withscores, score_cast_func, offset, num)
+
+    def zrevrange(self, name, start, end, withscores=False,
+                  score_cast_func=float):
+        """
+        Return a range of values from sorted set ``name`` between
+        ``start`` and ``end`` sorted in descending order.
+
+        ``start`` and ``end`` can be negative, indicating the end of the range.
+
+        ``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 desc:
-            return self.zrevrange(name, start, end, withscores,
-                                  score_cast_func)
-        pieces = ['ZRANGE', name, start, end]
-        if withscores:
-            pieces.append(b'WITHSCORES')
-        options = {
-            'withscores': withscores,
-            'score_cast_func': score_cast_func
-        }
-        return self.execute_command(*pieces, **options)
+        return self.zrange(name, start, end, desc=True,
+                           withscores=withscores,
+                           score_cast_func=score_cast_func)
 
-    def zrangestore(self, dest, name, start, end):
+    def zrangestore(self, dest, name, start, end,
+                    byscore=False, bylex=False, desc=False,
+                    offset=None, num=None):
         """
         Stores in ``dest`` the result of a range of values from sorted set
         ``name`` between ``start`` and ``end`` sorted in ascending order.
 
         ``start`` and ``end`` can be negative, indicating the end of the range.
+
+        ``byscore`` when set to True, returns the range of elements from the
+        sorted set having scores equal or between ``start`` and ``end``.
+
+        ``bylex`` when set to True, returns the range of elements from the
+        sorted set between the ``start`` and ``end`` lexicographical closed
+        range intervals.
+        Valid ``start`` and ``end`` must start with ( or [, in order to specify
+        whether the range interval is exclusive or inclusive, respectively.
+
+        ``desc`` a boolean indicating whether to sort the results in reversed
+        order.
+
+        ``offset`` and ``num`` are specified, then return a slice of the range.
+        Can't be provided when using ``bylex``.
         """
-        return self.execute_command('ZRANGESTORE', dest, name, start, end)
+        return self._zrange('ZRANGESTORE', dest, name, start, end, desc,
+                            byscore, bylex, False, None, offset, num)
 
     def zrangebylex(self, name, min, max, start=None, num=None):
         """
@@ -2484,13 +2550,7 @@ class Commands:
         If ``start`` and ``num`` are specified, then return a slice of the
         range.
         """
-        if (start is not None and num is None) or \
-                (num is not None and start is None):
-            raise DataError("``start`` and ``num`` must both be specified")
-        pieces = ['ZRANGEBYLEX', name, min, max]
-        if start is not None and num is not None:
-            pieces.extend([b'LIMIT', start, num])
-        return self.execute_command(*pieces)
+        return self.zrange(name, min, max, bylex=True, offset=start, num=num)
 
     def zrevrangebylex(self, name, max, min, start=None, num=None):
         """
@@ -2500,13 +2560,8 @@ class Commands:
         If ``start`` and ``num`` are specified, then return a slice of the
         range.
         """
-        if (start is not None and num is None) or \
-                (num is not None and start is None):
-            raise DataError("``start`` and ``num`` must both be specified")
-        pieces = ['ZREVRANGEBYLEX', name, max, min]
-        if start is not None and num is not None:
-            pieces.extend([b'LIMIT', start, num])
-        return self.execute_command(*pieces)
+        return self.zrange(name, max, min, desc=True,
+                           bylex=True, offset=start, num=num)
 
     def zrangebyscore(self, name, min, max, start=None, num=None,
                       withscores=False, score_cast_func=float):
@@ -2522,19 +2577,29 @@ class Commands:
 
         `score_cast_func`` a callable used to cast the score return value
         """
-        if (start is not None and num is None) or \
-                (num is not None and start is None):
-            raise DataError("``start`` and ``num`` must both be specified")
-        pieces = ['ZRANGEBYSCORE', name, min, max]
-        if start is not None and num is not None:
-            pieces.extend([b'LIMIT', start, num])
-        if withscores:
-            pieces.append(b'WITHSCORES')
-        options = {
-            'withscores': withscores,
-            'score_cast_func': score_cast_func
-        }
-        return self.execute_command(*pieces, **options)
+        return self.zrange(name, min, max, byscore=True,
+                           offset=start, num=num,
+                           withscores=withscores,
+                           score_cast_func=score_cast_func)
+
+    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.
+
+        If ``start`` and ``num`` are specified, then return a slice
+        of the range.
+
+        ``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, max, min, desc=True,
+                           byscore=True, offset=start,
+                           num=num, withscores=withscores,
+                           score_cast_func=score_cast_func)
 
     def zrank(self, name, value):
         """
@@ -2572,56 +2637,6 @@ class Commands:
         """
         return self.execute_command('ZREMRANGEBYSCORE', name, min, max)
 
-    def zrevrange(self, name, start, end, withscores=False,
-                  score_cast_func=float):
-        """
-        Return a range of values from sorted set ``name`` between
-        ``start`` and ``end`` sorted in descending order.
-
-        ``start`` and ``end`` can be negative, indicating the end of the range.
-
-        ``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
-        """
-        pieces = ['ZREVRANGE', name, start, end]
-        if withscores:
-            pieces.append(b'WITHSCORES')
-        options = {
-            'withscores': withscores,
-            'score_cast_func': score_cast_func
-        }
-        return self.execute_command(*pieces, **options)
-
-    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.
-
-        If ``start`` and ``num`` are specified, then return a slice
-        of the range.
-
-        ``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 (start is not None and num is None) or \
-                (num is not None and start is None):
-            raise DataError("``start`` and ``num`` must both be specified")
-        pieces = ['ZREVRANGEBYSCORE', name, max, min]
-        if start is not None and num is not None:
-            pieces.extend([b'LIMIT', start, num])
-        if withscores:
-            pieces.append(b'WITHSCORES')
-        options = {
-            'withscores': withscores,
-            'score_cast_func': score_cast_func
-        }
-        return self.execute_command(*pieces, **options)
-
     def zrevrank(self, name, value):
         """
         Returns a 0-based value indicating the descending rank of
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 904e27f..8929198 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -1840,6 +1840,7 @@ class TestRedisCommands:
         r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3})
         assert r.zrange('a', 0, 1) == [b'a1', b'a2']
         assert r.zrange('a', 1, 2) == [b'a2', b'a3']
+        assert r.zrange('a', 0, 2, desc=True) == [b'a3', b'a2', b'a1']
 
         # withscores
         assert r.zrange('a', 0, 1, withscores=True) == \
@@ -1851,6 +1852,46 @@ class TestRedisCommands:
         assert r.zrange('a', 0, 1, withscores=True, score_cast_func=int) == \
                [(b'a1', 1), (b'a2', 2)]
 
+    def test_zrange_errors(self, r):
+        with pytest.raises(exceptions.DataError):
+            r.zrange('a', 0, 1, byscore=True, bylex=True)
+        with pytest.raises(exceptions.DataError):
+            r.zrange('a', 0, 1, bylex=True, withscores=True)
+        with pytest.raises(exceptions.DataError):
+            r.zrange('a', 0, 1, byscore=True, withscores=True, offset=4)
+        with pytest.raises(exceptions.DataError):
+            r.zrange('a', 0, 1, byscore=True, withscores=True, num=2)
+
+    @skip_if_server_version_lt('6.2.0')
+    def test_zrange_params(self, r):
+        # bylex
+        r.zadd('a', {'a': 0, 'b': 0, 'c': 0, 'd': 0, 'e': 0, 'f': 0, 'g': 0})
+        assert r.zrange('a', '[aaa', '(g', bylex=True) == \
+               [b'b', b'c', b'd', b'e', b'f']
+        assert r.zrange('a', '[f', '+', bylex=True) == [b'f', b'g']
+        assert r.zrange('a', '+', '[f', desc=True, bylex=True) == [b'g', b'f']
+        assert r.zrange('a', '-', '+', bylex=True, offset=3, num=2) == \
+               [b'd', b'e']
+        assert r.zrange('a', '+', '-', desc=True, bylex=True,
+                        offset=3, num=2) == \
+               [b'd', b'c']
+
+        # byscore
+        r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3, 'a4': 4, 'a5': 5})
+        assert r.zrange('a', 2, 4, byscore=True, offset=1, num=2) == \
+               [b'a3', b'a4']
+        assert r.zrange('a', 4, 2, desc=True, byscore=True,
+                        offset=1, num=2) == \
+               [b'a3', b'a2']
+        assert r.zrange('a', 2, 4, byscore=True, withscores=True) == \
+               [(b'a2', 2.0), (b'a3', 3.0), (b'a4', 4.0)]
+        assert r.zrange('a', 4, 2, desc=True, byscore=True,
+                        withscores=True, score_cast_func=int) == \
+               [(b'a4', 4), (b'a3', 3), (b'a2', 2)]
+
+        # rev
+        assert r.zrange('a', 0, 1, desc=True) == [b'a5', b'a4']
+
     @skip_if_server_version_lt('6.2.0')
     def test_zrangestore(self, r):
         r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3})
@@ -1860,6 +1901,16 @@ class TestRedisCommands:
         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', 1, 2, byscore=True,
+                             offset=0, num=1)
+        assert r.zrange('b', 0, -1) == [b'a1']
+        # by lex
+        assert r.zrange('a', '[a2', '(a3', bylex=True) == \
+               [b'a2']
 
     @skip_if_server_version_lt('2.8.9')
     def test_zrangebylex(self, r):
@@ -1885,16 +1936,12 @@ class TestRedisCommands:
     def test_zrangebyscore(self, r):
         r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3, 'a4': 4, 'a5': 5})
         assert r.zrangebyscore('a', 2, 4) == [b'a2', b'a3', b'a4']
-
         # slicing with start/num
         assert r.zrangebyscore('a', 2, 4, start=1, num=2) == \
                [b'a3', b'a4']
-
         # withscores
         assert r.zrangebyscore('a', 2, 4, withscores=True) == \
                [(b'a2', 2.0), (b'a3', 3.0), (b'a4', 4.0)]
-
-        # custom score function
         assert r.zrangebyscore('a', 2, 4, withscores=True,
                                score_cast_func=int) == \
                [(b'a2', 2), (b'a3', 3), (b'a4', 4)]
@@ -1958,15 +2005,12 @@ class TestRedisCommands:
     def test_zrevrangebyscore(self, r):
         r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3, 'a4': 4, 'a5': 5})
         assert r.zrevrangebyscore('a', 4, 2) == [b'a4', b'a3', b'a2']
-
         # slicing with start/num
         assert r.zrevrangebyscore('a', 4, 2, start=1, num=2) == \
                [b'a3', b'a2']
-
         # withscores
         assert r.zrevrangebyscore('a', 4, 2, withscores=True) == \
                [(b'a4', 4.0), (b'a3', 3.0), (b'a2', 2.0)]
-
         # custom score function
         assert r.zrevrangebyscore('a', 4, 2, withscores=True,
                                   score_cast_func=int) == \