diff --git a/redis/client.py b/redis/client.py index e426abe0b5..c3282c78f0 100755 --- a/redis/client.py +++ b/redis/client.py @@ -148,6 +148,8 @@ def get_value(value): 'voted-leader-epoch': int } +GEOUNITS = ["m", "km", "mi", "ft"] + def parse_sentinel_state(item): result = pairs_to_dict_typed(item, SENTINEL_STATE_TYPES) @@ -1972,6 +1974,162 @@ def register_script(self, script): """ return Script(self, script) + def geoadd(self, key, *args): + """ + Adds georeferenced point under key (sorted set internaly). + + Arguments are any number of paramaters (divisible by 3): + longitude latitude member [longitude latitude member ...] + or single iterable (list) of: + [(longiture, latitude, member), ...] + + Returns number of elements added under key (serted set) + excluding already existing elements. + + Examples: + geoadd("Sicily", + 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + + geoadd("Sicily", [ + (13.361389, 38.115556, "Palermo"), + (15.087269, 37.502669, "Catania") + ]) + """ + pieces = [] + if len(args) == 1 and hasattr(args[0], '__iter__'): + for x in args[0]: + if hasattr(x, '__iter__'): + pieces.extend(x) + else: + pieces.append(x) + + elif len(args) % 3 != 0: + raise RedisError("invalid number of GEOADD arguments.") + else: + pieces = args + return self.execute_command('GEOADD', key, *pieces) + + def geodist(self, key, member1, member2, units='m'): + """ + Returns distance between member1 and member2 under key(sorted set) + in specified units: "m", "km", "mi", "ft" (default "m") + + Examples: + geodist("Sicily", "Palermo", "Catania", "km") + """ + return self.execute_command('GEODIST', key, member1, member2, + self._validate_units(units)) + + def geohash(self, key, *args): + """ + Returns valid Geohash string representing position + of one or more elements. + + Arguments are any number of members: + member [member ...] + or single iterable (list) of: + [member, ...] + + Examples: + geohash("Sicily", "Palermo", "Catania") + geohash("Sicily", ["Palermo", "Catania"]) + """ + pieces = [] + if len(args) == 1 and hasattr(args, '__iter__'): + pieces.extend(args) + else: + pieces = args + return self.execute_command('GEOHASH', key, *pieces) + + def geopos(self, key, *args): + """ + Returns the position (latitude, longitude) of one or more elements. + + Arguments are any number of members: + member [member ...] + or single iterable (list) of: + [member, ...] + + Examples: + geopos("Sicily", "Palermo", "Catania") + geopos("Sicily", ["Palermo", "Catania"] ) + """ + pieces = [] + if len(args) == 1 and hasattr(args, '__iter__'): + pieces.extend(args) + else: + pieces = args + return self.execute_command('GEOPOS', key, *pieces) + + def georadius(self, key, longitude, latitude, radius, units, + withcoord=None, withdist=None, withhash=None, count=None): + """ + Return the members of a key (sorted set) that are within range + of radius (center set by longitude, latitude) + + ``withcoord`` sets flag that will include coordinates of members. + + ``withdist`` sets flag that will include distance from center. + + ``withhash`` sets flag that will include 52bit int raw geohash. + + ``withhash`` sets limit of returned members. + """ + pieces = [longitude, latitude, radius, self._validate_units(units)] + pieces.extend(self._withgeoradius(withcoord, withdist, + withhash, count)) + + return self.execute_command('GEORADIUS', key, *pieces) + + def georadiusbymember(self, key, member, radius, units, + withcoord=None, withdist=None, + withhash=None, count=None): + """ + Return the members of a key (sorted set) that are within range + of radius (center set by member) + + ``withcoord`` sets flag that will include coordinates of members. + + ``withdist`` sets flag that will include distance from center. + + ``withhash`` sets flag that will include 52bit int raw geohash. + + ``withhash`` sets limit of returned members. + """ + pieces = [member, radius, self._validate_units(units)] + pieces.extend(self._withgeoradius(withcoord, withdist, + withhash, count)) + + return self.execute_command('GEORADIUSBYMEMBER', key, *pieces) + + def _withgeoradius(self, withcoord=None, withdist=None, + withhash=None, count=None): + pieces = [] + if withcoord: + pieces.append('WITHCOORD') + + if withdist: + pieces.append('WITHDIST') + + if withhash: + pieces.append('WITHHASH') + + if count is not None: + try: + int(count) + pieces.append('COUNT') + pieces.append(count) + except TypeError: + # count is not a number - ignore it + pass + return pieces + + def _validate_units(self, units): + if units not in GEOUNITS: + units = "m" + return units + class Redis(StrictRedis): """ diff --git a/tests/test_commands.py b/tests/test_commands.py index 7293810321..4caa16ce9b 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1338,6 +1338,114 @@ def test_sort_all_options(self, r): assert r.lrange('sorted', 0, 10) == \ [b('vodka'), b('milk'), b('gin'), b('apple juice')] + @skip_if_server_version_lt('3.2.0') + def test_geoadd(self, r): + first = r.geoadd("Sicily", 13.361389, 38.115556, + "Palermo", 15.087269, 37.502669, "Catania") + second = r.geoadd("Sicily", [13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania"]) + third = r.geoadd("Sicily", [(13.361389, 38.115556, "Palermo"), + (15.087269, 37.502669, "Catania")]) + assert first == 2 + assert second == 0 + assert third == 0 + + @skip_if_server_version_lt('3.2.0') + def test_geodist(self, r): + r.geoadd("Sicily", 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + dist1 = r.geodist("Sicily", "Palermo", "Catania") + dist2 = r.geodist("Sicily", "Palermo", "Catania", "km") + dist3 = r.geodist("Sicily", "Palermo", "Catania", "mi") + dist4 = r.geodist("Sicily", "Foo", "Bar") + assert abs(float(dist1) - float("166274.15156960039")) < 0.0000000001 + assert abs(float(dist2) - float("166.27415156960038")) < 0.0000000001 + assert abs(float(dist3) - float("103.31822459492736")) < 0.0000000001 + assert dist4 is None + + @skip_if_server_version_lt('3.2.0') + def test_geohash(self, r): + r.geoadd("Sicily", 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + hash = r.geohash("Sicily", "Palermo", "Catania") + assert set(hash) == set(["sqc8b49rny0", "sqdtr74hyu0"]) + assert set(hash) == set(["sqdtr74hyu0", "sqc8b49rny0"]) + + @skip_if_server_version_lt('3.2.0') + def test_geopos(self, r): + r.geoadd("Sicily", 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + pos = r.geopos("Sicily", "Palermo", "Catania", "NonExisting") + assert len(pos) == 3 + + assert len(pos[0]) == 2 + assert abs(float(pos[0][0]) - 13.361389) < 0.00001 + assert abs(float(pos[0][1]) - 38.115556) < 0.00001 + + assert len(pos[1]) == 2 + assert abs(float(pos[1][0]) - 15.087269) < 0.00001 + assert abs(float(pos[1][1]) - 37.502669) < 0.00001 + + assert pos[2] is None + + @skip_if_server_version_lt('3.2.0') + def test_georadius(self, r): + r.geoadd("Sicily", 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + radius1 = r.georadius("Sicily", 15, 37, 100, "km") + radius2 = r.georadius("Sicily", 15, 37, 200, "km") + assert radius1 == ['Catania'] + assert set(radius2) == set(["Catania", "Palermo"]) + assert set(radius2) == set(["Palermo", "Catania"]) + + extraradius1 = r.georadius("Sicily", 15, 37, 200, "km", withdist=True) + assert len(extraradius1) == 2 + assert len(extraradius1[0]) == 2 + assert len(extraradius1[1]) == 2 + + assert extraradius1[0][0] == "Palermo" + assert abs(float(extraradius1[0][1]) - 190.4424) < 0.01 + + assert extraradius1[1][0] == "Catania" + assert abs(float(extraradius1[1][1]) - 56.4413) < 0.01 + + extraradius2 = r.georadius("Sicily", 15, 37, 200, "km", withcoord=True) + assert len(extraradius2) == 2 + assert len(extraradius2[0]) == 2 + assert len(extraradius2[1]) == 2 + + assert extraradius2[0][0] == "Palermo" + assert len(extraradius2[0][1]) == 2 + assert abs(float(extraradius2[0][1][0]) - 13.361389) < 0.00001 + assert abs(float(extraradius2[0][1][1]) - 38.115556) < 0.00001 + + assert extraradius2[1][0] == "Catania" + assert abs(float(extraradius2[1][1][0]) - 15.087269) < 0.00001 + assert abs(float(extraradius2[1][1][1]) - 37.502669) < 0.00001 + + extraradius3 = r.georadius("Sicily", 15, 37, 200, "km", withhash=True) + assert len(extraradius3) == 2 + assert len(extraradius3[0]) == 2 + assert len(extraradius3[1]) == 2 + + assert extraradius3[0][0] == "Palermo" + assert extraradius3[0][1] == 3479099956230698 + + assert extraradius3[1][0] == "Catania" + assert extraradius3[1][1] == 3479447370796909 + + extraradius4 = r.georadius("Sicily", 15, 37, 200, "km", count=1) + assert len(extraradius4) == 1 + + @skip_if_server_version_lt('3.2.0') + def test_georadiusbymember(self, r): + r.geoadd("Sicily", 13.583333, 37.316667, "Agrigento") + r.geoadd("Sicily", 13.361389, 38.115556, "Palermo", + 15.087269, 37.502669, "Catania") + member_radius1 = r.georadiusbymember("Sicily", "Agrigento", 100, "km") + assert set(member_radius1) == set(["Agrigento", "Palermo"]) + assert set(member_radius1) == set(["Palermo", "Agrigento"]) + class TestStrictCommands(object):