Skip to content

Commit 03a0c31

Browse files
AvitalFineRedisAvitalFineRedis
authored
Add support to NX XX and CH to GEOADD (#1605)
1 parent 46430c2 commit 03a0c31

File tree

2 files changed

+77
-29
lines changed

2 files changed

+77
-29
lines changed

‎redis/commands.py‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2904,17 +2904,39 @@ def register_script(self, script):
29042904
return Script(self, script)
29052905

29062906
# GEO COMMANDS
2907-
def geoadd(self, name, *values):
2907+
def geoadd(self, name, values, nx=False, xx=False, ch=False):
29082908
"""
29092909
Add the specified geospatial items to the specified key identified
29102910
by the ``name`` argument. The Geospatial items are given as ordered
29112911
members of the ``values`` argument, each item or place is formed by
29122912
the triad longitude, latitude and name.
2913+
2914+
Note: You can use ZREM to remove elements.
2915+
2916+
``nx`` forces ZADD to only create new elements and not to update
2917+
scores for elements that already exist.
2918+
2919+
``xx`` forces ZADD to only update scores of elements that already
2920+
exist. New elements will not be added.
2921+
2922+
``ch`` modifies the return value to be the numbers of elements changed.
2923+
Changed elements include new elements that were added and elements
2924+
whose scores changed.
29132925
"""
2926+
if nx and xx:
2927+
raise DataError("GEOADD allows either 'nx' or 'xx', not both")
29142928
if len(values) % 3 != 0:
29152929
raise DataError("GEOADD requires places with lon, lat and name"
29162930
" values")
2917-
return self.execute_command('GEOADD', name, *values)
2931+
pieces = [name]
2932+
if nx:
2933+
pieces.append('NX')
2934+
if xx:
2935+
pieces.append('XX')
2936+
if ch:
2937+
pieces.append('CH')
2938+
pieces.extend(values)
2939+
return self.execute_command('GEOADD', *pieces)
29182940

29192941
def geodist(self, name, place1, place2, unit=None):
29202942
"""

‎tests/test_commands.py‎

Lines changed: 53 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2422,35 +2422,63 @@ def test_readonly(self, mock_cluster_resp_ok):
24222422
def test_geoadd(self, r):
24232423
values = (2.1909389952632, 41.433791470673, 'place1') + \
24242424
(2.1873744593677, 41.406342043777, 'place2')
2425-
2426-
assert r.geoadd('barcelona', *values) == 2
2425+
assert r.geoadd('barcelona', values) == 2
24272426
assert r.zcard('barcelona') == 2
24282427

2428+
@skip_if_server_version_lt('6.2.0')
2429+
def test_geoadd_nx(self, r):
2430+
values = (2.1909389952632, 41.433791470673, 'place1') + \
2431+
(2.1873744593677, 41.406342043777, 'place2')
2432+
assert r.geoadd('a', values) == 2
2433+
values = (2.1909389952632, 41.433791470673, 'place1') + \
2434+
(2.1873744593677, 41.406342043777, 'place2') + \
2435+
(2.1804738294738, 41.405647879212, 'place3')
2436+
assert r.geoadd('a', values, nx=True) == 1
2437+
assert r.zrange('a', 0, -1) == [b'place3', b'place2', b'place1']
2438+
2439+
@skip_if_server_version_lt('6.2.0')
2440+
def test_geoadd_xx(self, r):
2441+
values = (2.1909389952632, 41.433791470673, 'place1')
2442+
assert r.geoadd('a', values) == 1
2443+
values = (2.1909389952632, 41.433791470673, 'place1') + \
2444+
(2.1873744593677, 41.406342043777, 'place2')
2445+
assert r.geoadd('a', values, xx=True) == 0
2446+
assert r.zrange('a', 0, -1) == \
2447+
[b'place1']
2448+
2449+
@skip_if_server_version_lt('6.2.0')
2450+
def test_geoadd_ch(self, r):
2451+
values = (2.1909389952632, 41.433791470673, 'place1')
2452+
assert r.geoadd('a', values) == 1
2453+
values = (2.1909389952632, 31.433791470673, 'place1') + \
2454+
(2.1873744593677, 41.406342043777, 'place2')
2455+
assert r.geoadd('a', values, ch=True) == 2
2456+
assert r.zrange('a', 0, -1) == \
2457+
[b'place1', b'place2']
2458+
24292459
@skip_if_server_version_lt('3.2.0')
24302460
def test_geoadd_invalid_params(self, r):
24312461
with pytest.raises(exceptions.RedisError):
2432-
r.geoadd('barcelona', *(1, 2))
2462+
r.geoadd('barcelona', (1, 2))
24332463

24342464
@skip_if_server_version_lt('3.2.0')
24352465
def test_geodist(self, r):
24362466
values = (2.1909389952632, 41.433791470673, 'place1') + \
24372467
(2.1873744593677, 41.406342043777, 'place2')
2438-
2439-
assert r.geoadd('barcelona', *values) == 2
2468+
assert r.geoadd('barcelona', values) == 2
24402469
assert r.geodist('barcelona', 'place1', 'place2') == 3067.4157
24412470

24422471
@skip_if_server_version_lt('3.2.0')
24432472
def test_geodist_units(self, r):
24442473
values = (2.1909389952632, 41.433791470673, 'place1') + \
24452474
(2.1873744593677, 41.406342043777, 'place2')
2446-
2447-
r.geoadd('barcelona', *values)
2475+
r.geoadd('barcelona', values)
24482476
assert r.geodist('barcelona', 'place1', 'place2', 'km') == 3.0674
24492477

24502478
@skip_if_server_version_lt('3.2.0')
24512479
def test_geodist_missing_one_member(self, r):
24522480
values = (2.1909389952632, 41.433791470673, 'place1')
2453-
r.geoadd('barcelona', *values)
2481+
r.geoadd('barcelona', values)
24542482
assert r.geodist('barcelona', 'place1', 'missing_member', 'km') is None
24552483

24562484
@skip_if_server_version_lt('3.2.0')
@@ -2462,8 +2490,7 @@ def test_geodist_invalid_units(self, r):
24622490
def test_geohash(self, r):
24632491
values = (2.1909389952632, 41.433791470673, 'place1') + \
24642492
(2.1873744593677, 41.406342043777, 'place2')
2465-
2466-
r.geoadd('barcelona', *values)
2493+
r.geoadd('barcelona', values)
24672494
assert r.geohash('barcelona', 'place1', 'place2', 'place3') == \
24682495
['sp3e9yg3kd0', 'sp3e9cbc3t0', None]
24692496

@@ -2472,8 +2499,7 @@ def test_geohash(self, r):
24722499
def test_geopos(self, r):
24732500
values = (2.1909389952632, 41.433791470673, 'place1') + \
24742501
(2.1873744593677, 41.406342043777, 'place2')
2475-
2476-
r.geoadd('barcelona', *values)
2502+
r.geoadd('barcelona', values)
24772503
# redis uses 52 bits precision, hereby small errors may be introduced.
24782504
assert r.geopos('barcelona', 'place1', 'place2') == \
24792505
[(2.19093829393386841, 41.43379028184083523),
@@ -2493,7 +2519,7 @@ def test_geosearch(self, r):
24932519
values = (2.1909389952632, 41.433791470673, 'place1') + \
24942520
(2.1873744593677, 41.406342043777, b'\x80place2') + \
24952521
(2.583333, 41.316667, 'place3')
2496-
r.geoadd('barcelona', *values)
2522+
r.geoadd('barcelona', values)
24972523
assert r.geosearch('barcelona', longitude=2.191,
24982524
latitude=41.433, radius=1000) == [b'place1']
24992525
assert r.geosearch('barcelona', longitude=2.187,
@@ -2515,7 +2541,7 @@ def test_geosearch_member(self, r):
25152541
values = (2.1909389952632, 41.433791470673, 'place1') + \
25162542
(2.1873744593677, 41.406342043777, b'\x80place2')
25172543

2518-
r.geoadd('barcelona', *values)
2544+
r.geoadd('barcelona', values)
25192545
assert r.geosearch('barcelona', member='place1', radius=4000) == \
25202546
[b'\x80place2', b'place1']
25212547
assert r.geosearch('barcelona', member='place1', radius=10) == \
@@ -2534,7 +2560,7 @@ def test_geosearch_member(self, r):
25342560
def test_geosearch_sort(self, r):
25352561
values = (2.1909389952632, 41.433791470673, 'place1') + \
25362562
(2.1873744593677, 41.406342043777, 'place2')
2537-
r.geoadd('barcelona', *values)
2563+
r.geoadd('barcelona', values)
25382564
assert r.geosearch('barcelona', longitude=2.191,
25392565
latitude=41.433, radius=3000, sort='ASC') == \
25402566
[b'place1', b'place2']
@@ -2547,7 +2573,7 @@ def test_geosearch_sort(self, r):
25472573
def test_geosearch_with(self, r):
25482574
values = (2.1909389952632, 41.433791470673, 'place1') + \
25492575
(2.1873744593677, 41.406342043777, 'place2')
2550-
r.geoadd('barcelona', *values)
2576+
r.geoadd('barcelona', values)
25512577

25522578
# test a bunch of combinations to test the parse response
25532579
# function.
@@ -2618,7 +2644,7 @@ def test_geosearchstore(self, r):
26182644
values = (2.1909389952632, 41.433791470673, 'place1') + \
26192645
(2.1873744593677, 41.406342043777, 'place2')
26202646

2621-
r.geoadd('barcelona', *values)
2647+
r.geoadd('barcelona', values)
26222648
r.geosearchstore('places_barcelona', 'barcelona',
26232649
longitude=2.191, latitude=41.433, radius=1000)
26242650
assert r.zrange('places_barcelona', 0, -1) == [b'place1']
@@ -2629,7 +2655,7 @@ def test_geosearchstore_dist(self, r):
26292655
values = (2.1909389952632, 41.433791470673, 'place1') + \
26302656
(2.1873744593677, 41.406342043777, 'place2')
26312657

2632-
r.geoadd('barcelona', *values)
2658+
r.geoadd('barcelona', values)
26332659
r.geosearchstore('places_barcelona', 'barcelona',
26342660
longitude=2.191, latitude=41.433,
26352661
radius=1000, storedist=True)
@@ -2641,7 +2667,7 @@ def test_georadius(self, r):
26412667
values = (2.1909389952632, 41.433791470673, 'place1') + \
26422668
(2.1873744593677, 41.406342043777, b'\x80place2')
26432669

2644-
r.geoadd('barcelona', *values)
2670+
r.geoadd('barcelona', values)
26452671
assert r.georadius('barcelona', 2.191, 41.433, 1000) == [b'place1']
26462672
assert r.georadius('barcelona', 2.187, 41.406, 1000) == [b'\x80place2']
26472673

@@ -2650,15 +2676,15 @@ def test_georadius_no_values(self, r):
26502676
values = (2.1909389952632, 41.433791470673, 'place1') + \
26512677
(2.1873744593677, 41.406342043777, 'place2')
26522678

2653-
r.geoadd('barcelona', *values)
2679+
r.geoadd('barcelona', values)
26542680
assert r.georadius('barcelona', 1, 2, 1000) == []
26552681

26562682
@skip_if_server_version_lt('3.2.0')
26572683
def test_georadius_units(self, r):
26582684
values = (2.1909389952632, 41.433791470673, 'place1') + \
26592685
(2.1873744593677, 41.406342043777, 'place2')
26602686

2661-
r.geoadd('barcelona', *values)
2687+
r.geoadd('barcelona', values)
26622688
assert r.georadius('barcelona', 2.191, 41.433, 1, unit='km') == \
26632689
[b'place1']
26642690

@@ -2668,7 +2694,7 @@ def test_georadius_with(self, r):
26682694
values = (2.1909389952632, 41.433791470673, 'place1') + \
26692695
(2.1873744593677, 41.406342043777, 'place2')
26702696

2671-
r.geoadd('barcelona', *values)
2697+
r.geoadd('barcelona', values)
26722698

26732699
# test a bunch of combinations to test the parse response
26742700
# function.
@@ -2696,7 +2722,7 @@ def test_georadius_count(self, r):
26962722
values = (2.1909389952632, 41.433791470673, 'place1') + \
26972723
(2.1873744593677, 41.406342043777, 'place2')
26982724

2699-
r.geoadd('barcelona', *values)
2725+
r.geoadd('barcelona', values)
27002726
assert r.georadius('barcelona', 2.191, 41.433, 3000, count=1) == \
27012727
[b'place1']
27022728
assert r.georadius('barcelona', 2.191, 41.433, 3000,
@@ -2708,7 +2734,7 @@ def test_georadius_sort(self, r):
27082734
values = (2.1909389952632, 41.433791470673, 'place1') + \
27092735
(2.1873744593677, 41.406342043777, 'place2')
27102736

2711-
r.geoadd('barcelona', *values)
2737+
r.geoadd('barcelona', values)
27122738
assert r.georadius('barcelona', 2.191, 41.433, 3000, sort='ASC') == \
27132739
[b'place1', b'place2']
27142740
assert r.georadius('barcelona', 2.191, 41.433, 3000, sort='DESC') == \
@@ -2719,7 +2745,7 @@ def test_georadius_store(self, r):
27192745
values = (2.1909389952632, 41.433791470673, 'place1') + \
27202746
(2.1873744593677, 41.406342043777, 'place2')
27212747

2722-
r.geoadd('barcelona', *values)
2748+
r.geoadd('barcelona', values)
27232749
r.georadius('barcelona', 2.191, 41.433, 1000, store='places_barcelona')
27242750
assert r.zrange('places_barcelona', 0, -1) == [b'place1']
27252751

@@ -2729,7 +2755,7 @@ def test_georadius_store_dist(self, r):
27292755
values = (2.1909389952632, 41.433791470673, 'place1') + \
27302756
(2.1873744593677, 41.406342043777, 'place2')
27312757

2732-
r.geoadd('barcelona', *values)
2758+
r.geoadd('barcelona', values)
27332759
r.georadius('barcelona', 2.191, 41.433, 1000,
27342760
store_dist='places_barcelona')
27352761
# instead of save the geo score, the distance is saved.
@@ -2741,7 +2767,7 @@ def test_georadiusmember(self, r):
27412767
values = (2.1909389952632, 41.433791470673, 'place1') + \
27422768
(2.1873744593677, 41.406342043777, b'\x80place2')
27432769

2744-
r.geoadd('barcelona', *values)
2770+
r.geoadd('barcelona', values)
27452771
assert r.georadiusbymember('barcelona', 'place1', 4000) == \
27462772
[b'\x80place2', b'place1']
27472773
assert r.georadiusbymember('barcelona', 'place1', 10) == [b'place1']

0 commit comments

Comments
 (0)