Skip to content

Commit

Permalink
added GeoLocationField with auto index-creation for GEO2D
Browse files Browse the repository at this point in the history
  • Loading branch information
flosch committed Mar 29, 2010
1 parent 38b2919 commit 2304dac
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 5 deletions.
19 changes: 18 additions & 1 deletion mongoengine/fields.py
Expand Up @@ -12,7 +12,7 @@
'DateTimeField', 'EmbeddedDocumentField', 'ListField', 'DictField',
'ObjectIdField', 'ReferenceField', 'ValidationError',
'DecimalField', 'URLField', 'GenericReferenceField',
'BinaryField', 'EmailField']
'BinaryField', 'EmailField', 'GeoLocationField']

RECURSIVE_REFERENCE_CONSTANT = 'self'

Expand Down Expand Up @@ -342,6 +342,23 @@ def validate(self, value):
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
Expand Down
12 changes: 8 additions & 4 deletions mongoengine/queryset.py
Expand Up @@ -4,7 +4,6 @@
import re
import copy


__all__ = ['queryset_manager', 'Q', 'InvalidQueryError',
'InvalidCollectionError']

Expand Down Expand Up @@ -228,7 +227,7 @@ def _collection(self):
"""
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']:
Expand All @@ -242,6 +241,11 @@ def _collection(self):
# 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
Expand Down Expand Up @@ -297,7 +301,7 @@ def _transform_query(cls, _doc_cls=None, **query):
"""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']

Expand All @@ -308,7 +312,7 @@ def _transform_query(cls, _doc_cls=None, **query):
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)
Expand Down

0 comments on commit 2304dac

Please sign in to comment.