From 2304dac8e3b92714b978e71d336cb2d66f248201 Mon Sep 17 00:00:00 2001 From: Florian Schlachter Date: Tue, 30 Mar 2010 00:04:39 +0200 Subject: [PATCH] added GeoLocationField with auto index-creation for GEO2D --- mongoengine/fields.py | 19 ++++++++++++++++++- mongoengine/queryset.py | 12 ++++++++---- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 622145e9..4c15dc0f 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', 'EmailField'] + 'BinaryField', 'EmailField', 'GeoLocationField'] RECURSIVE_REFERENCE_CONSTANT = 'self' @@ -342,6 +342,23 @@ class DictField(BaseField): def lookup_member(self, member_name): return BaseField(db_field=member_name) +class GeoLocationField(DictField): + """Supports geobased fields""" + + def validate(self, value): + """Make sure that a geo-value is of type (x, y) + """ + if not isinstance(value, tuple) and not isinstance(value, list): + raise ValidationError('GeoLocationField can only hold tuples or lists of (x, y)') + + if len(value) <> 2: + raise ValidationError('GeoLocationField must have exactly two elements (x, y)') + + def to_mongo(self, value): + return {'x': value[0], 'y': value[1]} + + def to_python(self, value): + return value.keys() class ReferenceField(BaseField): """A reference to a document that will be automatically dereferenced on diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 11dc2bcb..a3bae3ce 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -4,7 +4,6 @@ import pymongo import re import copy - __all__ = ['queryset_manager', 'Q', 'InvalidQueryError', 'InvalidCollectionError'] @@ -228,7 +227,7 @@ class QuerySet(object): """ if not self._accessed_collection: self._accessed_collection = True - + # Ensure document-defined indexes are created if self._document._meta['indexes']: for key_or_list in self._document._meta['indexes']: @@ -242,6 +241,11 @@ class QuerySet(object): # If _types is being used (for polymorphism), it needs an index if '_types' in self._query: self._collection.ensure_index('_types') + + # Ensure all needed field indexes are created + for field_name, field_instance in self._document._fields.iteritems(): + if field_instance.__class__.__name__ == 'GeoLocationField': + self._collection.ensure_index([(field_name, pymongo.GEO2D),]) return self._collection_obj @property @@ -297,7 +301,7 @@ class QuerySet(object): """Transform a query from Django-style format to Mongo format. """ operators = ['ne', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod', - 'all', 'size', 'exists'] + 'all', 'size', 'exists', 'near'] match_operators = ['contains', 'icontains', 'startswith', 'istartswith', 'endswith', 'iendswith'] @@ -308,7 +312,7 @@ class QuerySet(object): op = None if parts[-1] in operators + match_operators: op = parts.pop() - + if _doc_cls: # Switch field names to proper names [set in Field(name='foo')] fields = QuerySet._lookup_field(_doc_cls, parts)