Merge branch 'master' of git://github.com/flosch/mongoengine
Added unit test for get_or_create, merged flosch's get with punteney's get. Conflicts: mongoengine/queryset.py
This commit is contained in:
@@ -6,7 +6,6 @@ import pymongo
|
||||
class ValidationError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class BaseField(object):
|
||||
"""A base class for fields in a MongoDB document. Instances of this class
|
||||
may be added to subclasses of `Document` to define a document's schema.
|
||||
@@ -76,7 +75,10 @@ class ObjectIdField(BaseField):
|
||||
|
||||
def to_mongo(self, value):
|
||||
if not isinstance(value, pymongo.objectid.ObjectId):
|
||||
return pymongo.objectid.ObjectId(str(value))
|
||||
try:
|
||||
return pymongo.objectid.ObjectId(str(value))
|
||||
except Exception, e:
|
||||
raise ValidationError(e.message)
|
||||
return value
|
||||
|
||||
def prepare_query_value(self, value):
|
||||
|
||||
@@ -209,20 +209,19 @@ class ReferenceField(BaseField):
|
||||
return super(ReferenceField, self).__get__(instance, owner)
|
||||
|
||||
def to_mongo(self, document):
|
||||
if isinstance(document, (str, unicode, pymongo.objectid.ObjectId)):
|
||||
# document may already be an object id
|
||||
id_ = document
|
||||
else:
|
||||
id_field_name = self.document_type._meta['id_field']
|
||||
id_field = self.document_type._fields[id_field_name]
|
||||
|
||||
if isinstance(document, Document):
|
||||
# We need the id from the saved object to create the DBRef
|
||||
id_ = document.id
|
||||
if id_ is None:
|
||||
raise ValidationError('You can only reference documents once '
|
||||
'they have been saved to the database')
|
||||
else:
|
||||
id_ = document
|
||||
|
||||
# id may be a string rather than an ObjectID object
|
||||
if not isinstance(id_, pymongo.objectid.ObjectId):
|
||||
id_ = pymongo.objectid.ObjectId(id_)
|
||||
|
||||
id_ = id_field.to_mongo(id_)
|
||||
collection = self.document_type._meta['collection']
|
||||
return pymongo.dbref.DBRef(collection, id_)
|
||||
|
||||
|
||||
@@ -10,6 +10,13 @@ __all__ = ['queryset_manager', 'Q', 'InvalidQueryError',
|
||||
# The maximum number of items to display in a QuerySet.__repr__
|
||||
REPR_OUTPUT_SIZE = 20
|
||||
|
||||
class DoesNotExist(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class MultipleObjectsReturned(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidQueryError(Exception):
|
||||
pass
|
||||
@@ -298,18 +305,43 @@ class QuerySet(object):
|
||||
|
||||
def get(self, *q_objs, **query):
|
||||
"""Retrieve the the matching object raising
|
||||
'MultipleObjectsReturned' or 'DoesNotExist' exceptions
|
||||
if multiple or no results are found.
|
||||
:class:`~mongoengine.queryset.MultipleObjectsReturned` or
|
||||
:class:`~mongoengine.queryset.DoesNotExist` exceptions if multiple or
|
||||
no results are found.
|
||||
"""
|
||||
self.__call__(*q_objs, **query)
|
||||
count = self.count()
|
||||
if count == 1:
|
||||
return self[0]
|
||||
elif count > 1:
|
||||
raise MultipleObjectsReturned
|
||||
message = u'%d items returned, instead of 1' % count
|
||||
raise MultipleObjectsReturned(message)
|
||||
else:
|
||||
raise DoesNotExist
|
||||
raise DoesNotExist('Document not found')
|
||||
|
||||
def get_or_create(self, *q_objs, **query):
|
||||
"""Retreive unique object or create, if it doesn't exist. Raises
|
||||
:class:`~mongoengine.queryset.MultipleObjectsReturned` if multiple
|
||||
results are found. A new document will be created if the document
|
||||
doesn't exists; a dictionary of default values for the new document
|
||||
may be provided as a keyword argument called :attr:`defaults`.
|
||||
"""
|
||||
defaults = query.get('defaults', {})
|
||||
if query.has_key('defaults'):
|
||||
del query['defaults']
|
||||
|
||||
self.__call__(*q_objs, **query)
|
||||
count = self.count()
|
||||
if count == 0:
|
||||
query.update(defaults)
|
||||
doc = self._document(**query)
|
||||
doc.save()
|
||||
return doc
|
||||
elif count == 1:
|
||||
return self.first()
|
||||
else:
|
||||
message = u'%d items returned, instead of 1' % count
|
||||
raise MultipleObjectsReturned(message)
|
||||
|
||||
def first(self):
|
||||
"""Retrieve the first object matching the query.
|
||||
|
||||
Reference in New Issue
Block a user