From a4d2f22fd2016e47987eb8c48f541db99eafb067 Mon Sep 17 00:00:00 2001 From: Matt Dennewitz Date: Tue, 23 Mar 2010 00:14:01 -0500 Subject: [PATCH 1/2] added 'geo_indexes' to TopLevelDocumentMetaclass; added GeoPointField, a glorified [lat float, lng float] container; added geo lookup operators to QuerySet; added initial geo tests --- mongoengine/base.py | 1 + mongoengine/fields.py | 17 ++++++++++++- mongoengine/queryset.py | 26 ++++++++++++++++---- tests/queryset.py | 54 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 92 insertions(+), 6 deletions(-) diff --git a/mongoengine/base.py b/mongoengine/base.py index 0a904467..e4b45f26 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -206,6 +206,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass): 'max_size': None, 'ordering': [], # default ordering applied at runtime 'indexes': [], # indexes to be ensured at runtime + 'geo_indexes': [], 'id_field': id_field, } diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 9a9f4e0e..bb61f1da 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -12,7 +12,7 @@ __all__ = ['StringField', 'IntField', 'FloatField', 'BooleanField', 'DateTimeField', 'EmbeddedDocumentField', 'ListField', 'DictField', 'ObjectIdField', 'ReferenceField', 'ValidationError', 'DecimalField', 'URLField', 'GenericReferenceField', - 'BinaryField'] + 'BinaryField', 'GeoPointField'] RECURSIVE_REFERENCE_CONSTANT = 'self' @@ -443,6 +443,7 @@ class GenericReferenceField(BaseField): def prepare_query_value(self, op, value): return self.to_mongo(value)['_ref'] + class BinaryField(BaseField): """A binary data field. """ @@ -462,3 +463,17 @@ class BinaryField(BaseField): if self.max_bytes is not None and len(value) > self.max_bytes: raise ValidationError('Binary value is too long') + + +class GeoPointField(BaseField): + """A list storing a latitude and longitude. + """ + + def validate(self, value): + assert isinstance(value, (list, tuple)) + + if not len(value) == 2: + raise ValidationError('Value must be a two-dimensional point.') + if not isinstance(value[0], (float, int)) and \ + not isinstance(value[1], (float, int)): + raise ValidationError('Both values in point must be float or int.') diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 11dc2bcb..f03f97a5 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -232,12 +232,17 @@ class QuerySet(object): # Ensure document-defined indexes are created if self._document._meta['indexes']: for key_or_list in self._document._meta['indexes']: - #self.ensure_index(key_or_list) self._collection.ensure_index(key_or_list) # Ensure indexes created by uniqueness constraints for index in self._document._meta['unique_indexes']: self._collection.ensure_index(index, unique=True) + + if self._document._meta['geo_indexes'] and \ + pymongo.version >= "1.5.1": + from pymongo import GEO2D + for index in self._document._meta['geo_indexes']: + self._collection.ensure_index([(index, GEO2D)]) # If _types is being used (for polymorphism), it needs an index if '_types' in self._query: @@ -298,6 +303,7 @@ class QuerySet(object): """ operators = ['ne', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod', 'all', 'size', 'exists'] + geo_operators = ['within_distance', 'within_box', 'near'] match_operators = ['contains', 'icontains', 'startswith', 'istartswith', 'endswith', 'iendswith'] @@ -306,7 +312,7 @@ class QuerySet(object): parts = key.split('__') # Check for an operator and transform to mongo-style if there is op = None - if parts[-1] in operators + match_operators: + if parts[-1] in operators + match_operators + geo_operators: op = parts.pop() if _doc_cls: @@ -320,15 +326,25 @@ class QuerySet(object): singular_ops += match_operators if op in singular_ops: value = field.prepare_query_value(op, value) - elif op in ('in', 'nin', 'all'): + elif op in ('in', 'nin', 'all', 'near'): # 'in', 'nin' and 'all' require a list of values value = [field.prepare_query_value(op, v) for v in value] if field.__class__.__name__ == 'GenericReferenceField': parts.append('_ref') - if op and op not in match_operators: - value = {'$' + op: value} + # if op and op not in match_operators: + if op: + if op in geo_operators: + if op == "within_distance": + value = {'$within': {'$center': value}} + elif op == "near": + value = {'$near': value} + else: + raise NotImplmenetedError, \ + "Geo method has been implemented" + elif op not in match_operators: + value = {'$' + op: value} key = '.'.join(parts) if op is None or key not in mongo_query: diff --git a/tests/queryset.py b/tests/queryset.py index c0bd8fa9..004c4160 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1070,6 +1070,60 @@ class QuerySetTest(unittest.TestCase): def tearDown(self): self.Person.drop_collection() + def test_near(self): + """Ensure that "near" queries work with and without radii. + """ + class Event(Document): + title = StringField() + location = GeoPointField() + + def __unicode__(self): + return self.title + + meta = {'geo_indexes': ["location"]} + + Event.drop_collection() + + event1 = Event(title="Coltrane Motion @ Double Door", + location=[41.909889, -87.677137]) + event2 = Event(title="Coltrane Motion @ Bottom of the Hill", + location=[37.7749295, -122.4194155]) + event3 = Event(title="Coltrane Motion @ Empty Bottle", + location=[41.900474, -87.686638]) + + event1.save() + event2.save() + event3.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(location__near=[41.9120459, -87.67892]) + self.assertEqual(events.count(), 3) + self.assertEqual(list(events), [event1, event3, event2]) + + # find events within 5 miles of pitchfork office, chicago + point_and_distance = [[41.9120459, -87.67892], 5] + events = Event.objects(location__within_distance=point_and_distance) + self.assertEqual(events.count(), 2) + events = list(events) + self.assertTrue(event2 not in events) + self.assertTrue(event1 in events) + self.assertTrue(event3 in events) + + # find events around san francisco + point_and_distance = [[37.7566023, -122.415579], 10] + events = Event.objects(location__within_distance=point_and_distance) + self.assertEqual(events.count(), 1) + self.assertEqual(events[0], event2) + + # find events within 1 mile of greenpoint, broolyn, nyc, ny + point_and_distance = [[40.7237134, -73.9509714], 1] + events = Event.objects(location__within_distance=point_and_distance) + self.assertEqual(events.count(), 0) + + Event.drop_collection() + class QTest(unittest.TestCase): From 600ca3bcf9d9f87129f8504286f7e2108154b3cd Mon Sep 17 00:00:00 2001 From: Matt Dennewitz Date: Tue, 23 Mar 2010 00:57:26 -0500 Subject: [PATCH 2/2] renamed 'test_near' to 'test_geospatial_operators', updated added ordering checks to test --- tests/queryset.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/queryset.py b/tests/queryset.py index 004c4160..c8db29a2 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -1070,11 +1070,12 @@ class QuerySetTest(unittest.TestCase): def tearDown(self): self.Person.drop_collection() - def test_near(self): - """Ensure that "near" queries work with and without radii. + def test_geospatial_operators(self): + """Ensure that geospatial queries are working. """ class Event(Document): title = StringField() + date = DateTimeField() location = GeoPointField() def __unicode__(self): @@ -1085,10 +1086,13 @@ class QuerySetTest(unittest.TestCase): Event.drop_collection() event1 = Event(title="Coltrane Motion @ Double Door", + date=datetime.now() - timedelta(days=1), location=[41.909889, -87.677137]) event2 = Event(title="Coltrane Motion @ Bottom of the Hill", + date=datetime.now() - timedelta(days=10), location=[37.7749295, -122.4194155]) event3 = Event(title="Coltrane Motion @ Empty Bottle", + date=datetime.now(), location=[41.900474, -87.686638]) event1.save() @@ -1111,6 +1115,12 @@ class QuerySetTest(unittest.TestCase): self.assertTrue(event1 in events) self.assertTrue(event3 in events) + # ensure ordering is respected by "near" + events = Event.objects(location__near=[41.9120459, -87.67892]) + events = events.order_by("-date") + self.assertEqual(events.count(), 3) + self.assertEqual(list(events), [event3, event1, event2]) + # find events around san francisco point_and_distance = [[37.7566023, -122.415579], 10] events = Event.objects(location__within_distance=point_and_distance) @@ -1122,6 +1132,13 @@ class QuerySetTest(unittest.TestCase): events = Event.objects(location__within_distance=point_and_distance) self.assertEqual(events.count(), 0) + # ensure ordering is respected by "within_distance" + point_and_distance = [[41.9120459, -87.67892], 10] + events = Event.objects(location__within_distance=point_and_distance) + events = events.order_by("-date") + self.assertEqual(events.count(), 2) + self.assertEqual(events[0], event3) + Event.drop_collection()