Skip to content

Add support to NX XX and CH to GEOADD #1605

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Oct 18, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions redis/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -2865,17 +2865,39 @@ 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.

Note: You can use ZREM to remove elements.

``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):
"""
Expand Down
80 changes: 53 additions & 27 deletions tests/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -2378,35 +2378,63 @@ def test_readonly(self, mock_cluster_resp_ok):
def test_geoadd(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

assert r.geoadd('barcelona', *values) == 2
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):
r.geoadd('barcelona', *(1, 2))
r.geoadd('barcelona', (1, 2))

@skip_if_server_version_lt('3.2.0')
def test_geodist(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

assert r.geoadd('barcelona', *values) == 2
assert r.geoadd('barcelona', values) == 2
assert r.geodist('barcelona', 'place1', 'place2') == 3067.4157

@skip_if_server_version_lt('3.2.0')
def test_geodist_units(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geodist('barcelona', 'place1', 'place2', 'km') == 3.0674

@skip_if_server_version_lt('3.2.0')
def test_geodist_missing_one_member(self, r):
values = (2.1909389952632, 41.433791470673, 'place1')
r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geodist('barcelona', 'place1', 'missing_member', 'km') is None

@skip_if_server_version_lt('3.2.0')
Expand All @@ -2418,8 +2446,7 @@ def test_geodist_invalid_units(self, r):
def test_geohash(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geohash('barcelona', 'place1', 'place2', 'place3') == \
['sp3e9yg3kd0', 'sp3e9cbc3t0', None]

Expand All @@ -2428,8 +2455,7 @@ def test_geohash(self, r):
def test_geopos(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
# redis uses 52 bits precision, hereby small errors may be introduced.
assert r.geopos('barcelona', 'place1', 'place2') == \
[(2.19093829393386841, 41.43379028184083523),
Expand All @@ -2449,7 +2475,7 @@ def test_geosearch(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, b'\x80place2') + \
(2.583333, 41.316667, 'place3')
r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geosearch('barcelona', longitude=2.191,
latitude=41.433, radius=1000) == [b'place1']
assert r.geosearch('barcelona', longitude=2.187,
Expand All @@ -2471,7 +2497,7 @@ def test_geosearch_member(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, b'\x80place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geosearch('barcelona', member='place1', radius=4000) == \
[b'\x80place2', b'place1']
assert r.geosearch('barcelona', member='place1', radius=10) == \
Expand All @@ -2490,7 +2516,7 @@ def test_geosearch_member(self, r):
def test_geosearch_sort(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')
r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.geosearch('barcelona', longitude=2.191,
latitude=41.433, radius=3000, sort='ASC') == \
[b'place1', b'place2']
Expand All @@ -2503,7 +2529,7 @@ def test_geosearch_sort(self, r):
def test_geosearch_with(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')
r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)

# test a bunch of combinations to test the parse response
# function.
Expand Down Expand Up @@ -2574,7 +2600,7 @@ def test_geosearchstore(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
r.geosearchstore('places_barcelona', 'barcelona',
longitude=2.191, latitude=41.433, radius=1000)
assert r.zrange('places_barcelona', 0, -1) == [b'place1']
Expand All @@ -2585,7 +2611,7 @@ def test_geosearchstore_dist(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
r.geosearchstore('places_barcelona', 'barcelona',
longitude=2.191, latitude=41.433,
radius=1000, storedist=True)
Expand All @@ -2597,7 +2623,7 @@ def test_georadius(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, b'\x80place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadius('barcelona', 2.191, 41.433, 1000) == [b'place1']
assert r.georadius('barcelona', 2.187, 41.406, 1000) == [b'\x80place2']

Expand All @@ -2606,15 +2632,15 @@ def test_georadius_no_values(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadius('barcelona', 1, 2, 1000) == []

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

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadius('barcelona', 2.191, 41.433, 1, unit='km') == \
[b'place1']

Expand All @@ -2624,7 +2650,7 @@ def test_georadius_with(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)

# test a bunch of combinations to test the parse response
# function.
Expand Down Expand Up @@ -2652,7 +2678,7 @@ def test_georadius_count(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadius('barcelona', 2.191, 41.433, 3000, count=1) == \
[b'place1']

Expand All @@ -2661,7 +2687,7 @@ def test_georadius_sort(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadius('barcelona', 2.191, 41.433, 3000, sort='ASC') == \
[b'place1', b'place2']
assert r.georadius('barcelona', 2.191, 41.433, 3000, sort='DESC') == \
Expand All @@ -2672,7 +2698,7 @@ def test_georadius_store(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

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

Expand All @@ -2682,7 +2708,7 @@ def test_georadius_store_dist(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, 'place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
r.georadius('barcelona', 2.191, 41.433, 1000,
store_dist='places_barcelona')
# instead of save the geo score, the distance is saved.
Expand All @@ -2694,7 +2720,7 @@ def test_georadiusmember(self, r):
values = (2.1909389952632, 41.433791470673, 'place1') + \
(2.1873744593677, 41.406342043777, b'\x80place2')

r.geoadd('barcelona', *values)
r.geoadd('barcelona', values)
assert r.georadiusbymember('barcelona', 'place1', 4000) == \
[b'\x80place2', b'place1']
assert r.georadiusbymember('barcelona', 'place1', 10) == [b'place1']
Expand Down