Skip to content

Commit

Permalink
LT and GT support for ZADD (redis#1509)
Browse files Browse the repository at this point in the history
Co-authored-by: malinaa96 <52569986+malinaa96@users.noreply.github.com>
Co-authored-by: Avital Fine <79420960+AvitalFineRedis@users.noreply.github.com>
  • Loading branch information
3 people committed Jul 29, 2021
1 parent 1503c72 commit 75ad1cb
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 1 deletion.
10 changes: 9 additions & 1 deletion redis/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2871,7 +2871,8 @@ def xtrim(self, name, maxlen, approximate=True):
return self.execute_command('XTRIM', name, *pieces)

# SORTED SET COMMANDS
def zadd(self, name, mapping, nx=False, xx=False, ch=False, incr=False):
def zadd(self, name, mapping, nx=False, xx=False, ch=False, incr=False,
gt=None, lt=None):
"""
Set any number of element-name, score pairs to the key ``name``. Pairs
are specified as a dict of element-names keys to score values.
Expand Down Expand Up @@ -2902,6 +2903,9 @@ def zadd(self, name, mapping, nx=False, xx=False, ch=False, incr=False):
if incr and len(mapping) != 1:
raise DataError("ZADD option 'incr' only works when passing a "
"single element/score pair")
if nx is True and (gt is not None or lt is not None):
raise DataError("Only one of 'nx', 'lt', or 'gr' may be defined.")

pieces = []
options = {}
if nx:
Expand All @@ -2913,6 +2917,10 @@ def zadd(self, name, mapping, nx=False, xx=False, ch=False, incr=False):
if incr:
pieces.append(b'INCR')
options['as_score'] = True
if gt:
pieces.append(b'GT')
if lt:
pieces.append(b'LT')
for pair in mapping.items():
pieces.append(pair[1])
pieces.append(pair[0])
Expand Down
17 changes: 17 additions & 0 deletions tests/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -1470,6 +1470,23 @@ def test_zadd_incr_with_xx(self, r):
# redis-py
assert r.zadd('a', {'a1': 1}, xx=True, incr=True) is None

@skip_if_server_version_lt('6.2.0')
def test_zadd_gt_lt(self, r):

for i in range(1, 20):
r.zadd('a', {'a%s' % i: i})
assert r.zadd('a', {'a20': 5}, gt=3) == 1

for i in range(1, 20):
r.zadd('a', {'a%s' % i: i})
assert r.zadd('a', {'a2': 5}, lt=1) == 0

# cannot use both nx and xx options
with pytest.raises(exceptions.DataError):
r.zadd('a', {'a15': 155}, nx=True, lt=True)
r.zadd('a', {'a15': 155}, nx=True, gt=True)
r.zadd('a', {'a15': 155}, lx=True, gt=True)

def test_zcard(self, r):
r.zadd('a', {'a1': 1, 'a2': 2, 'a3': 3})
assert r.zcard('a') == 3
Expand Down

0 comments on commit 75ad1cb

Please sign in to comment.