LT and GT support for ZADD (#1509)

Co-authored-by: malinaa96 <[email protected]>
Co-authored-by: Avital Fine <[email protected]>

Author: 3 people

redis/client.py

@@ -2866,7 +2866,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.
@@ -2897,6 +2898,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:
@@ -2908,6 +2912,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])

tests/test_commands.py

@@ -1462,6 +1462,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