Add support to ANY to GEOSEARCHSTORE and to GEOSEARCH

redis · Oct 13, 2021 · 9046884 · 9046884
1 parent 5cd878d
commit 9046884
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git 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
@@ -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):