Skip to content

Commit

Permalink
Implement GEO commands
Browse files Browse the repository at this point in the history
  • Loading branch information
cunla committed Feb 24, 2023
1 parent ed00059 commit 12daf86
Show file tree
Hide file tree
Showing 4 changed files with 65 additions and 12 deletions.
18 changes: 9 additions & 9 deletions docs/redis-commands/Redis.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
7 changes: 5 additions & 2 deletions fakeredis/_command_args_parsing.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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 ~.
Expand Down Expand Up @@ -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:
Expand Down
16 changes: 15 additions & 1 deletion fakeredis/commands_mixins/geo_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)
36 changes: 36 additions & 0 deletions test/test_mixins/test_geo_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",))
Expand Down Expand Up @@ -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"]

0 comments on commit 12daf86

Please sign in to comment.