Implement GEO commands

Author: cunla

docs/redis-commands/Redis.md

@@ -1198,30 +1198,30 @@ Returns members of a geospatial index as standard geohash strings
 
 Returns longitude and latitude of members of a geospatial index
 
-
-### Unsupported geo commands 
-> To implement support for a command, see [here](/guides/implement-command/) 
-
-#### [GEORADIUS](https://redis.io/commands/georadius/) <small>(not implemented)</small>
+### [GEORADIUS](https://redis.io/commands/georadius/)
 
 Query a sorted set representing a geospatial index to fetch members matching a given maximum distance from a point
 
-#### [GEORADIUSBYMEMBER](https://redis.io/commands/georadiusbymember/) <small>(not implemented)</small>
+### [GEORADIUSBYMEMBER](https://redis.io/commands/georadiusbymember/)
 
 Query a sorted set representing a geospatial index to fetch members matching a given maximum distance from a member
 
-#### [GEORADIUSBYMEMBER_RO](https://redis.io/commands/georadiusbymember_ro/) <small>(not implemented)</small>
+### [GEORADIUSBYMEMBER_RO](https://redis.io/commands/georadiusbymember_ro/)
 
 A read-only variant for GEORADIUSBYMEMBER
 
-#### [GEORADIUS_RO](https://redis.io/commands/georadius_ro/) <small>(not implemented)</small>
+### [GEORADIUS_RO](https://redis.io/commands/georadius_ro/)
 
 A read-only variant for GEORADIUS
 
-#### [GEOSEARCH](https://redis.io/commands/geosearch/) <small>(not implemented)</small>
+### [GEOSEARCH](https://redis.io/commands/geosearch/)
 
 Query a sorted set representing a geospatial index to fetch members inside an area of a box or a circle.
 
+
+### Unsupported geo commands 
+> To implement support for a command, see [here](/guides/implement-command/) 
+
 #### [GEOSEARCHSTORE](https://redis.io/commands/geosearchstore/) <small>(not implemented)</small>
 
 Query a sorted set representing a geospatial index to fetch members inside an area of a box or a circle, and store the result in another key.

fakeredis/_command_args_parsing.py

@@ -1,13 +1,13 @@
 from typing import Tuple, List, Dict, Any
 
 from . import _msgs as msgs
-from ._commands import Int
+from ._commands import Int, Float
 from ._helpers import SimpleError, null_terminate
 
 
 def _count_params(s: str):
     res = 0
-    while s[res] in '+*~':
+    while s[res] in '.+*~':
         res += 1
     return res
 
@@ -54,6 +54,7 @@ def extract_args(
 
     An expected argument can have parameters:
     - A numerical (Int) parameter is identified with +.
+    - A float (Float) parameter is identified with .
     - A non-numerical parameter is identified with a *.
     - A argument with potentially ~ or = between the
       argument name and the value is identified with a ~.
@@ -110,6 +111,8 @@ def _parse_params(
             curr_arg = actual_args[ind + i + 1]
             if argument_name[i] == '+':
                 curr_arg = Int.decode(curr_arg)
+            elif argument_name[i] == '.':
+                curr_arg = Float.decode(curr_arg)
             temp_res.append(curr_arg)
 
         if len(temp_res) == 1:

fakeredis/commands_mixins/geo_mixin.py

@@ -142,7 +142,7 @@ def geodist(self, key, m1, m2, *args):
 
     def _search(
             self, key, long, lat, radius, conv,
-            withcoord, withdist, withhash, count, count_any, desc, store, storedist):
+            withcoord, withdist, _, count, count_any, desc, store, storedist):
         zset = key.value
         geo_results = _find_near(zset, lat, long, radius, conv, count, count_any, desc)
 
@@ -188,3 +188,17 @@ def georadiusbymember_ro(self, key, member_name, radius, *args):
         member_score = key.value.get(member_name)
         lat, long, _, _ = geohash.decode(member_score)
         return self.georadius_ro(key, long, lat, radius, *args)
+
+    @command(name='GEOSEARCH', fixed=(Key(ZSet),), repeat=(bytes,))
+    def geosearch(self, key, *args):
+        (frommember, (long, lat), radius), left_args = extract_args(
+            args, ('*frommember', '..fromlonlat', '.byradius'),
+            error_on_unexpected=False, left_from_first_unexpected=False)
+        if frommember is None and long is None:
+            raise SimpleError(msgs.SYNTAX_ERROR_MSG)
+        if frommember is not None and long is not None:
+            raise SimpleError(msgs.SYNTAX_ERROR_MSG)
+        if frommember:
+            return self.georadiusbymember_ro(key, frommember, radius, *left_args)
+        else:
+            return self.georadius_ro(key, long, lat, radius, *left_args)

test/test_mixins/test_geo_commands.py

@@ -120,6 +120,25 @@ def test_georadius(
     assert r.georadius("barcelona", long, lat, radius, **extra) == expected
 
 
+@pytest.mark.parametrize(
+    "member,radius,extra,expected", [
+        ('place1', 1000, {}, [b"place1"]),
+        ('place2', 1000, {}, [b"place2"]),
+        ('place1', 1, {"unit": "km"}, [b"place1"]),
+        ('place1', 3000, {"count": 1}, [b"place1"]),
+    ])
+def test_georadiusbymember(
+        r: redis.Redis, member: str, radius: float,
+        extra: Dict[str, Any],
+        expected):
+    values = ((2.1909389952632, 41.433791470673, "place1") +
+              (2.1873744593677, 41.406342043777, b"place2"))
+    r.geoadd("barcelona", values)
+    assert r.georadiusbymember("barcelona", member, radius, **extra) == expected
+    assert r.georadiusbymember("barcelona", member, radius, **extra, store_dist='extract') == len(expected)
+    assert r.zcard("extract") == len(expected)
+
+
 def test_georadius_with(r: redis.Redis):
     values = ((2.1909389952632, 41.433791470673, "place1") +
               (2.1873744593677, 41.406342043777, "place2",))
@@ -180,3 +199,20 @@ def test_georadius_errors(r: redis.Redis):
         r.geoadd('newgroup', bad_values)
     with pytest.raises(redis.ResponseError):
         testtools.raw_command(r, 'geoadd', 'newgroup', *bad_values)
+
+
+def test_geosearch(r: redis.Redis):
+    values = (
+            (2.1909389952632, 41.433791470673, "place1")
+            + (2.1873744593677, 41.406342043777, b"place2")
+            + (2.583333, 41.316667, "place3")
+    )
+    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, latitude=41.406, radius=1000) == [b"place2"]
+    # assert r.geosearch("barcelona", longitude=2.191, latitude=41.433, height=1000, width=1000) == [b"place1"]
+    assert set(r.geosearch("barcelona", member="place3", radius=100, unit="km")) == {b"place2", b"place1", b"place3", }
+    # test count
+    assert r.geosearch("barcelona", member="place3", radius=100, unit="km", count=2) == [b"place3", b"place2"]
+    assert r.geosearch("barcelona", member="place3", radius=100, unit="km", count=1, any=1)[0] in [
+        b"place1", b"place3", b"place2"]