From c7ecb1d4aa68820e000f88c23f6ff3afd4166b8a Mon Sep 17 00:00:00 2001
From: Chayim <chayim@users.noreply.github.com>
Date: Thu, 29 Jul 2021 14:46:50 +0300
Subject: LT and GT support for ZADD (#1509)

Co-authored-by: malinaa96 <52569986+malinaa96@users.noreply.github.com>
Co-authored-by: Avital Fine <79420960+AvitalFineRedis@users.noreply.github.com>
---
 redis/client.py        | 10 +++++++++-
 tests/test_commands.py | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/redis/client.py b/redis/client.py
index 160f495..ea36ef0 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -2866,7 +2866,8 @@ class Redis:
         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.
@@ -2897,6 +2898,9 @@ class Redis:
         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:
@@ -2908,6 +2912,10 @@ class Redis:
         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])
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 3f0a82f..62394a4 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -1462,6 +1462,23 @@ class TestRedisCommands:
         # 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
-- 
cgit v1.2.1