From 90468846acb7812df6419d54a9c191156e170516 Mon Sep 17 00:00:00 2001 From: AvitalFineRedis Date: Wed, 13 Oct 2021 15:33:49 +0200 Subject: [PATCH] Add support to ANY to GEOSEARCHSTORE and to GEOSEARCH --- redis/commands.py | 31 +++++++++++++++++++------------ tests/test_commands.py | 6 ++++++ 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/redis/commands.py b/redis/commands.py index eb7cea54f6..ecf5b5137a 100644 --- a/redis/commands.py +++ b/redis/commands.py @@ -2908,7 +2908,7 @@ def geopos(self, name, *values): def georadius(self, name, longitude, latitude, radius, unit=None, withdist=False, withcoord=False, withhash=False, count=None, - sort=None, store=None, store_dist=None): + sort=None, store=None, store_dist=None, any=False): """ Return the members of the specified key identified by the ``name`` argument which are within the borders of the area specified @@ -2942,11 +2942,12 @@ def georadius(self, name, longitude, latitude, radius, unit=None, unit=unit, withdist=withdist, withcoord=withcoord, withhash=withhash, count=count, sort=sort, store=store, - store_dist=store_dist) + store_dist=store_dist, any=any) def georadiusbymember(self, name, member, radius, unit=None, withdist=False, withcoord=False, withhash=False, - count=None, sort=None, store=None, store_dist=None): + count=None, sort=None, store=None, store_dist=None, + any=False): """ This command is exactly like ``georadius`` with the sole difference that instead of taking, as the center of the area to query, a longitude @@ -2958,7 +2959,7 @@ def georadiusbymember(self, name, member, radius, unit=None, withdist=withdist, withcoord=withcoord, withhash=withhash, count=count, sort=sort, store=store, - store_dist=store_dist) + store_dist=store_dist, any=any) def _georadiusgeneric(self, command, *args, **kwargs): pieces = list(args) @@ -2969,21 +2970,26 @@ def _georadiusgeneric(self, command, *args, **kwargs): else: pieces.append('m',) + if kwargs['any'] and kwargs['count'] is None: + raise DataError("``any`` can't be provided without ``count``") + for arg_name, byte_repr in ( - ('withdist', b'WITHDIST'), - ('withcoord', b'WITHCOORD'), - ('withhash', b'WITHHASH')): + ('withdist', 'WITHDIST'), + ('withcoord', 'WITHCOORD'), + ('withhash', 'WITHHASH')): if kwargs[arg_name]: pieces.append(byte_repr) - if kwargs['count']: - pieces.extend([b'COUNT', kwargs['count']]) + if kwargs['count'] is not None: + pieces.extend(['COUNT', kwargs['count']]) + if kwargs['any']: + pieces.append('ANY') if kwargs['sort']: if kwargs['sort'] == 'ASC': - pieces.append(b'ASC') + pieces.append('ASC') elif kwargs['sort'] == 'DESC': - pieces.append(b'DESC') + pieces.append('DESC') else: raise DataError("GEORADIUS invalid sort") @@ -3116,7 +3122,8 @@ def _geosearchgeneric(self, command, *args, **kwargs): if kwargs['any']: pieces.append(b'ANY') elif kwargs['any']: - raise DataError("GEOSEARCH any can't be provided without count") + raise DataError("GEOSEARCH ``any`` can't be provided " + "without count") # other properties for arg_name, byte_repr in ( diff --git a/tests/test_commands.py b/tests/test_commands.py index 2be8923e0e..df79940913 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -2655,6 +2655,9 @@ def test_georadius_count(self, r): r.geoadd('barcelona', *values) assert r.georadius('barcelona', 2.191, 41.433, 3000, count=1) == \ [b'place1'] + assert r.georadius('barcelona', 2.191, 41.433, 3000, + count=1, any=True) == \ + [b'place2'] @skip_if_server_version_lt('3.2.0') def test_georadius_sort(self, r): @@ -2706,6 +2709,9 @@ def test_georadiusmember(self, r): (2.187376320362091, 41.40634178640635)], [b'place1', 0.0, 3471609698139488, (2.1909382939338684, 41.433790281840835)]] + assert r.georadiusbymember('barcelona', 'place1', 4000, + count=1, any=True) == \ + [b'\x80place2'] @skip_if_server_version_lt('5.0.0') def test_xack(self, r):