Add support to NX XX and CH to GEOADD

redis/commands.py

@@ -2865,17 +2865,37 @@ def register_script(self, script):
         return Script(self, script)
 
     # GEO COMMANDS
-    def geoadd(self, name, *values):
+    def geoadd(self, name, *values, nx=False, xx=False, ch=False):
         """
         Add the specified geospatial items to the specified key identified
         by the ``name`` argument. The Geospatial items are given as ordered
         members of the ``values`` argument, each item or place is formed by
         the triad longitude, latitude and name.
+
+        ``nx`` forces ZADD to only create new elements and not to update
+        scores for elements that already exist.
+
+        ``xx`` forces ZADD to only update scores of elements that already
+        exist. New elements will not be added.
+
+        ``ch`` modifies the return value to be the numbers of elements changed.
+        Changed elements include new elements that were added and elements
+        whose scores changed.
         """
+        if nx and xx:
+            raise DataError("GEOADD allows either 'nx' or 'xx', not both")
         if len(values) % 3 != 0:
             raise DataError("GEOADD requires places with lon, lat and name"
                             " values")
-        return self.execute_command('GEOADD', name, *values)
+        pieces = [name]
+        if nx:
+            pieces.append('NX')
+        if xx:
+            pieces.append('XX')
+        if ch:
+            pieces.append('CH')
+        pieces.extend(values)
+        return self.execute_command('GEOADD', *pieces)
 
     def geodist(self, name, place1, place2, unit=None):
         """

tests/test_commands.py

@@ -2382,6 +2382,37 @@ def test_geoadd(self, r):
         assert r.geoadd('barcelona', *values) == 2
         assert r.zcard('barcelona') == 2
 
+    @skip_if_server_version_lt('6.2.0')
+    def test_geoadd_nx(self, r):
+        values = (2.1909389952632, 41.433791470673, 'place1') + \
+                 (2.1873744593677, 41.406342043777, 'place2')
+        assert r.geoadd('a', *values) == 2
+        values = (2.1909389952632, 41.433791470673, 'place1') + \
+                 (2.1873744593677, 41.406342043777, 'place2') + \
+                 (2.1804738294738, 41.405647879212, 'place3')
+        assert r.geoadd('a', *values, nx=True) == 1
+        assert r.zrange('a', 0, -1) == [b'place3', b'place2', b'place1']
+
+    @skip_if_server_version_lt('6.2.0')
+    def test_geoadd_xx(self, r):
+        values = (2.1909389952632, 41.433791470673, 'place1')
+        assert r.geoadd('a', *values) == 1
+        values = (2.1909389952632, 41.433791470673, 'place1') + \
+                 (2.1873744593677, 41.406342043777, 'place2')
+        assert r.geoadd('a', *values, xx=True) == 0
+        assert r.zrange('a', 0, -1) == \
+               [b'place1']
+
+    @skip_if_server_version_lt('6.2.0')
+    def test_geoadd_ch(self, r):
+        values = (2.1909389952632, 41.433791470673, 'place1')
+        assert r.geoadd('a', *values) == 1
+        values = (2.1909389952632, 31.433791470673, 'place1') + \
+                 (2.1873744593677, 41.406342043777, 'place2')
+        assert r.geoadd('a', *values, ch=True) == 2
+        assert r.zrange('a', 0, -1) == \
+               [b'place1', b'place2']
+
     @skip_if_server_version_lt('3.2.0')
     def test_geoadd_invalid_params(self, r):
         with pytest.raises(exceptions.RedisError):