pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added 'geo_indexes' to TopLevelDocumentMetaclass; added GeoPointField, a glorified [lat float, lng float] container; added geo lookup operators to QuerySet; added initial geo tests

Browse Source
This commit is contained in:
Matt Dennewitz 2010-03-23 00:14:01 -05:00
parent 00c8d7e6f5
commit a4d2f22fd2
4 changed files with 92 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mongoengine/base.py
View File

@ -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,
}

17
mongoengine/fields.py
View File

@ -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.')

26
mongoengine/queryset.py
View File

@ -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:

54
tests/queryset.py
View File

@ -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):