import datetime import unittest from mongoengine import * from tests.utils import MongoDBTestCase, requires_mongodb_gte_3 __all__ = ("GeoQueriesTest",) class GeoQueriesTest(MongoDBTestCase): def _create_event_data(self, point_field_class=GeoPointField): """Create some sample data re-used in many of the tests below.""" class Event(Document): title = StringField() date = DateTimeField() location = point_field_class() def __unicode__(self): return self.title self.Event = Event Event.drop_collection() event1 = Event.objects.create( title="Coltrane Motion @ Double Door", date=datetime.datetime.now() - datetime.timedelta(days=1), location=[-87.677137, 41.909889]) event2 = Event.objects.create( title="Coltrane Motion @ Bottom of the Hill", date=datetime.datetime.now() - datetime.timedelta(days=10), location=[-122.4194155, 37.7749295]) event3 = Event.objects.create( title="Coltrane Motion @ Empty Bottle", date=datetime.datetime.now(), location=[-87.686638, 41.900474]) return event1, event2, event3 def test_near(self): """Make sure the "near" operator works.""" event1, event2, event3 = self._create_event_data() # find all events "near" pitchfork office, chicago. # note that "near" will show the san francisco event, too, # although it sorts to last. events = self.Event.objects(location__near=[-87.67892, 41.9120459]) self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event1, event3, event2]) # ensure ordering is respected by "near" events = self.Event.objects(location__near=[-87.67892, 41.9120459]) events = events.order_by("-date") self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event3, event1, event2]) def test_near_and_max_distance(self): """Ensure the "max_distance" operator works alongside the "near" operator. """ event1, event2, event3 = self._create_event_data() # find events within 10 degrees of san francisco point = [-122.415579, 37.7566023] events = self.Event.objects(location__near=point, location__max_distance=10) self.assertEqual(events.count(), 1) self.assertEqual(events[0], event2) # $minDistance was added in MongoDB v2.6, but continued being buggy # until v3.0; skip for older versions @requires_mongodb_gte_3 def test_near_and_min_distance(self): """Ensure the "min_distance" operator works alongside the "near" operator. """ event1, event2, event3 = self._create_event_data() # find events at least 10 degrees away of san francisco point = [-122.415579, 37.7566023] events = self.Event.objects(location__near=point, location__min_distance=10) self.assertEqual(events.count(), 2) def test_within_distance(self): """Make sure the "within_distance" operator works.""" event1, event2, event3 = self._create_event_data() # find events within 5 degrees of pitchfork office, chicago point_and_distance = [[-87.67892, 41.9120459], 5] events = self.Event.objects( location__within_distance=point_and_distance) self.assertEqual(events.count(), 2) events = list(events) self.assertNotIn(event2, events) self.assertIn(event1, events) self.assertIn(event3, events) # find events within 10 degrees of san francisco point_and_distance = [[-122.415579, 37.7566023], 10] events = self.Event.objects( location__within_distance=point_and_distance) self.assertEqual(events.count(), 1) self.assertEqual(events[0], event2) # find events within 1 degree of greenpoint, broolyn, nyc, ny point_and_distance = [[-73.9509714, 40.7237134], 1] events = self.Event.objects( location__within_distance=point_and_distance) self.assertEqual(events.count(), 0) # ensure ordering is respected by "within_distance" point_and_distance = [[-87.67892, 41.9120459], 10] events = self.Event.objects( location__within_distance=point_and_distance) events = events.order_by("-date") self.assertEqual(events.count(), 2) self.assertEqual(events[0], event3) def test_within_box(self): """Ensure the "within_box" operator works.""" event1, event2, event3 = self._create_event_data() # check that within_box works box = [(-125.0, 35.0), (-100.0, 40.0)] events = self.Event.objects(location__within_box=box) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event2.id) def test_within_polygon(self): """Ensure the "within_polygon" operator works.""" event1, event2, event3 = self._create_event_data() polygon = [ (-87.694445, 41.912114), (-87.69084, 41.919395), (-87.681742, 41.927186), (-87.654276, 41.911731), (-87.656164, 41.898061), ] events = self.Event.objects(location__within_polygon=polygon) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event1.id) polygon2 = [ (-1.742249, 54.033586), (-1.225891, 52.792797), (-4.40094, 53.389881) ] events = self.Event.objects(location__within_polygon=polygon2) self.assertEqual(events.count(), 0) def test_2dsphere_near(self): """Make sure the "near" operator works with a PointField, which corresponds to a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) # find all events "near" pitchfork office, chicago. # note that "near" will show the san francisco event, too, # although it sorts to last. events = self.Event.objects(location__near=[-87.67892, 41.9120459]) self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event1, event3, event2]) # ensure ordering is respected by "near" events = self.Event.objects(location__near=[-87.67892, 41.9120459]) events = events.order_by("-date") self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event3, event1, event2]) def test_2dsphere_near_and_max_distance(self): """Ensure the "max_distance" operator works alongside the "near" operator with a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) # find events within 10km of san francisco point = [-122.415579, 37.7566023] events = self.Event.objects(location__near=point, location__max_distance=10000) self.assertEqual(events.count(), 1) self.assertEqual(events[0], event2) # find events within 1km of greenpoint, broolyn, nyc, ny events = self.Event.objects(location__near=[-73.9509714, 40.7237134], location__max_distance=1000) self.assertEqual(events.count(), 0) # ensure ordering is respected by "near" events = self.Event.objects( location__near=[-87.67892, 41.9120459], location__max_distance=10000 ).order_by("-date") self.assertEqual(events.count(), 2) self.assertEqual(events[0], event3) def test_2dsphere_geo_within_box(self): """Ensure the "geo_within_box" operator works with a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) # check that within_box works box = [(-125.0, 35.0), (-100.0, 40.0)] events = self.Event.objects(location__geo_within_box=box) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event2.id) def test_2dsphere_geo_within_polygon(self): """Ensure the "geo_within_polygon" operator works with a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) polygon = [ (-87.694445, 41.912114), (-87.69084, 41.919395), (-87.681742, 41.927186), (-87.654276, 41.911731), (-87.656164, 41.898061), ] events = self.Event.objects(location__geo_within_polygon=polygon) self.assertEqual(events.count(), 1) self.assertEqual(events[0].id, event1.id) polygon2 = [ (-1.742249, 54.033586), (-1.225891, 52.792797), (-4.40094, 53.389881) ] events = self.Event.objects(location__geo_within_polygon=polygon2) self.assertEqual(events.count(), 0) # $minDistance was added in MongoDB v2.6, but continued being buggy # until v3.0; skip for older versions @requires_mongodb_gte_3 def test_2dsphere_near_and_min_max_distance(self): """Ensure "min_distace" and "max_distance" operators work well together with the "near" operator in a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) # ensure min_distance and max_distance combine well events = self.Event.objects( location__near=[-87.67892, 41.9120459], location__min_distance=1000, location__max_distance=10000 ).order_by("-date") self.assertEqual(events.count(), 1) self.assertEqual(events[0], event3) # ensure ordering is respected by "near" with "min_distance" events = self.Event.objects( location__near=[-87.67892, 41.9120459], location__min_distance=10000 ).order_by("-date") self.assertEqual(events.count(), 1) self.assertEqual(events[0], event2) def test_2dsphere_geo_within_center(self): """Make sure the "geo_within_center" operator works with a 2dsphere index. """ event1, event2, event3 = self._create_event_data( point_field_class=PointField ) # find events within 5 degrees of pitchfork office, chicago point_and_distance = [[-87.67892, 41.9120459], 2] events = self.Event.objects( location__geo_within_center=point_and_distance) self.assertEqual(events.count(), 2) events = list(events) self.assertNotIn(event2, events) self.assertIn(event1, events) self.assertIn(event3, events) def _test_embedded(self, point_field_class): """Helper test method ensuring given point field class works well in an embedded document. """ class Venue(EmbeddedDocument): location = point_field_class() name = StringField() class Event(Document): title = StringField() venue = EmbeddedDocumentField(Venue) Event.drop_collection() venue1 = Venue(name="The Rock", location=[-87.677137, 41.909889]) venue2 = Venue(name="The Bridge", location=[-122.4194155, 37.7749295]) event1 = Event(title="Coltrane Motion @ Double Door", venue=venue1).save() event2 = Event(title="Coltrane Motion @ Bottom of the Hill", venue=venue2).save() event3 = Event(title="Coltrane Motion @ Empty Bottle", venue=venue1).save() # find all events "near" pitchfork office, chicago. # note that "near" will show the san francisco event, too, # although it sorts to last. events = Event.objects(venue__location__near=[-87.67892, 41.9120459]) self.assertEqual(events.count(), 3) self.assertEqual(list(events), [event1, event3, event2]) def test_geo_spatial_embedded(self): """Make sure GeoPointField works properly in an embedded document.""" self._test_embedded(point_field_class=GeoPointField) def test_2dsphere_point_embedded(self): """Make sure PointField works properly in an embedded document.""" self._test_embedded(point_field_class=PointField) # Needs MongoDB > 2.6.4 https://jira.mongodb.org/browse/SERVER-14039 @requires_mongodb_gte_3 def test_spherical_geospatial_operators(self): """Ensure that spherical geospatial queries are working.""" class Point(Document): location = GeoPointField() Point.drop_collection() # These points are one degree apart, which (according to Google Maps) # is about 110 km apart at this place on the Earth. north_point = Point(location=[-122, 38]).save() # Near Concord, CA south_point = Point(location=[-122, 37]).save() # Near Santa Cruz, CA earth_radius = 6378.009 # in km (needs to be a float for dividing by) # Finds both points because they are within 60 km of the reference # point equidistant between them. points = Point.objects(location__near_sphere=[-122, 37.5]) self.assertEqual(points.count(), 2) # Same behavior for _within_spherical_distance points = Point.objects( location__within_spherical_distance=[ [-122, 37.5], 60 / earth_radius ] ) self.assertEqual(points.count(), 2) points = Point.objects(location__near_sphere=[-122, 37.5], location__max_distance=60 / earth_radius) self.assertEqual(points.count(), 2) # Test query works with max_distance, being farer from one point points = Point.objects(location__near_sphere=[-122, 37.8], location__max_distance=60 / earth_radius) close_point = points.first() self.assertEqual(points.count(), 1) # Test query works with min_distance, being farer from one point points = Point.objects(location__near_sphere=[-122, 37.8], location__min_distance=60 / earth_radius) self.assertEqual(points.count(), 1) far_point = points.first() self.assertNotEqual(close_point, far_point) # Finds both points, but orders the north point first because it's # closer to the reference point to the north. points = Point.objects(location__near_sphere=[-122, 38.5]) self.assertEqual(points.count(), 2) self.assertEqual(points[0].id, north_point.id) self.assertEqual(points[1].id, south_point.id) # Finds both points, but orders the south point first because it's # closer to the reference point to the south. points = Point.objects(location__near_sphere=[-122, 36.5]) self.assertEqual(points.count(), 2) self.assertEqual(points[0].id, south_point.id) self.assertEqual(points[1].id, north_point.id) # Finds only one point because only the first point is within 60km of # the reference point to the south. points = Point.objects( location__within_spherical_distance=[ [-122, 36.5], 60 / earth_radius ] ) self.assertEqual(points.count(), 1) self.assertEqual(points[0].id, south_point.id) def test_linestring(self): class Road(Document): name = StringField() line = LineStringField() Road.drop_collection() Road(name="66", line=[[40, 5], [41, 6]]).save() # near point = {"type": "Point", "coordinates": [40, 5]} roads = Road.objects.filter(line__near=point["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__near=point).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__near={"$geometry": point}).count() self.assertEqual(1, roads) # Within polygon = {"type": "Polygon", "coordinates": [[[40, 5], [40, 6], [41, 6], [41, 5], [40, 5]]]} roads = Road.objects.filter(line__geo_within=polygon["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_within=polygon).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_within={"$geometry": polygon}).count() self.assertEqual(1, roads) # Intersects line = {"type": "LineString", "coordinates": [[40, 5], [40, 6]]} roads = Road.objects.filter(line__geo_intersects=line["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_intersects=line).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_intersects={"$geometry": line}).count() self.assertEqual(1, roads) polygon = {"type": "Polygon", "coordinates": [[[40, 5], [40, 6], [41, 6], [41, 5], [40, 5]]]} roads = Road.objects.filter(line__geo_intersects=polygon["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_intersects=polygon).count() self.assertEqual(1, roads) roads = Road.objects.filter(line__geo_intersects={"$geometry": polygon}).count() self.assertEqual(1, roads) def test_polygon(self): class Road(Document): name = StringField() poly = PolygonField() Road.drop_collection() Road(name="66", poly=[[[40, 5], [40, 6], [41, 6], [40, 5]]]).save() # near point = {"type": "Point", "coordinates": [40, 5]} roads = Road.objects.filter(poly__near=point["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__near=point).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__near={"$geometry": point}).count() self.assertEqual(1, roads) # Within polygon = {"type": "Polygon", "coordinates": [[[40, 5], [40, 6], [41, 6], [41, 5], [40, 5]]]} roads = Road.objects.filter(poly__geo_within=polygon["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_within=polygon).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_within={"$geometry": polygon}).count() self.assertEqual(1, roads) # Intersects line = {"type": "LineString", "coordinates": [[40, 5], [41, 6]]} roads = Road.objects.filter(poly__geo_intersects=line["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_intersects=line).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_intersects={"$geometry": line}).count() self.assertEqual(1, roads) polygon = {"type": "Polygon", "coordinates": [[[40, 5], [40, 6], [41, 6], [41, 5], [40, 5]]]} roads = Road.objects.filter(poly__geo_intersects=polygon["coordinates"]).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_intersects=polygon).count() self.assertEqual(1, roads) roads = Road.objects.filter(poly__geo_intersects={"$geometry": polygon}).count() self.assertEqual(1, roads) def test_aspymongo_with_only(self): """Ensure as_pymongo works with only""" class Place(Document): location = PointField() Place.drop_collection() p = Place(location=[24.946861267089844, 60.16311983618494]) p.save() qs = Place.objects().only('location') self.assertDictEqual( qs.as_pymongo()[0]['location'], {u'type': u'Point', u'coordinates': [ 24.946861267089844, 60.16311983618494] } ) def test_2dsphere_point_sets_correctly(self): class Location(Document): loc = PointField() Location.drop_collection() Location(loc=[1,2]).save() loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["loc"], {"type": "Point", "coordinates": [1, 2]}) Location.objects.update(set__loc=[2,1]) loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["loc"], {"type": "Point", "coordinates": [2, 1]}) def test_2dsphere_linestring_sets_correctly(self): class Location(Document): line = LineStringField() Location.drop_collection() Location(line=[[1, 2], [2, 2]]).save() loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["line"], {"type": "LineString", "coordinates": [[1, 2], [2, 2]]}) Location.objects.update(set__line=[[2, 1], [1, 2]]) loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["line"], {"type": "LineString", "coordinates": [[2, 1], [1, 2]]}) def test_geojson_PolygonField(self): class Location(Document): poly = PolygonField() Location.drop_collection() Location(poly=[[[40, 5], [40, 6], [41, 6], [40, 5]]]).save() loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["poly"], {"type": "Polygon", "coordinates": [[[40, 5], [40, 6], [41, 6], [40, 5]]]}) Location.objects.update(set__poly=[[[40, 4], [40, 6], [41, 6], [40, 4]]]) loc = Location.objects.as_pymongo()[0] self.assertEqual(loc["poly"], {"type": "Polygon", "coordinates": [[[40, 4], [40, 6], [41, 6], [40, 4]]]}) if __name__ == '__main__': unittest.main()