diff --git a/mongoengine/base.py b/mongoengine/base.py index 0a904467..fbd95b3c 100644 --- a/mongoengine/base.py +++ b/mongoengine/base.py @@ -1,7 +1,9 @@ -from queryset import QuerySet, QuerySetManager - +import sys import pymongo +from queryset import QuerySet, QuerySetManager +from queryset import DoesNotExist, MultipleObjectsReturned + _document_registry = {} @@ -167,7 +169,22 @@ class DocumentMetaclass(type): for field in new_class._fields.values(): field.owner_document = new_class + module = attrs.pop('__module__') + + new_class.add_to_class('DoesNotExist', subclass_exception('DoesNotExist', + tuple(x.DoesNotExist + for k,x in superclasses.items()) + or (DoesNotExist,), module)) + + new_class.add_to_class('MultipleObjectsReturned', subclass_exception('MultipleObjectsReturned', + tuple(x.MultipleObjectsReturned + for k,x in superclasses.items()) + or (MultipleObjectsReturned,), module)) return new_class + + + def add_to_class(self, name, value): + setattr(self, name, value) class TopLevelDocumentMetaclass(DocumentMetaclass): @@ -417,3 +434,11 @@ class BaseDocument(object): if self.id == other.id: return True return False + +if sys.version_info < (2, 5): + # Prior to Python 2.5, Exception was an old-style class + def subclass_exception(name, parents, unused): + return types.ClassType(name, parents, {}) +else: + def subclass_exception(name, parents, module): + return type(name, parents, {'__module__': module}) \ No newline at end of file diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 11dc2bcb..0623719d 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -352,9 +352,10 @@ class QuerySet(object): return self[0] elif count > 1: message = u'%d items returned, instead of 1' % count - raise MultipleObjectsReturned(message) + raise self._document.MultipleObjectsReturned(message) else: - raise DoesNotExist('Document not found') + raise self._document.DoesNotExist("%s matching query does not exist." + % self._document._class_name) def get_or_create(self, *q_objs, **query): """Retreive unique object or create, if it doesn't exist. Raises @@ -380,7 +381,7 @@ class QuerySet(object): return self.first() else: message = u'%d items returned, instead of 1' % count - raise MultipleObjectsReturned(message) + raise self._document.MultipleObjectsReturned(message) def first(self): """Retrieve the first object matching the query. diff --git a/tests/queryset.py b/tests/queryset.py index c0bd8fa9..668f0355 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -147,6 +147,7 @@ class QuerySetTest(unittest.TestCase): """ # Try retrieving when no objects exists self.assertRaises(DoesNotExist, self.Person.objects.get) + self.assertRaises(self.Person.DoesNotExist, self.Person.objects.get) person1 = self.Person(name="User A", age=20) person1.save() @@ -155,6 +156,7 @@ class QuerySetTest(unittest.TestCase): # Retrieve the first person from the database self.assertRaises(MultipleObjectsReturned, self.Person.objects.get) + self.assertRaises(self.Person.MultipleObjectsReturned, self.Person.objects.get) # Use a query to filter the people found to just person2 person = self.Person.objects.get(age=30) @@ -162,6 +164,9 @@ class QuerySetTest(unittest.TestCase): person = self.Person.objects.get(age__lt=30) self.assertEqual(person.name, "User A") + + + def test_get_or_create(self): """Ensure that ``get_or_create`` returns one result or creates a new @@ -175,6 +180,8 @@ class QuerySetTest(unittest.TestCase): # Retrieve the first person from the database self.assertRaises(MultipleObjectsReturned, self.Person.objects.get_or_create) + self.assertRaises(self.Person.MultipleObjectsReturned, + self.Person.objects.get_or_create) # Use a query to filter the people found to just person2 person = self.Person.objects.get_or_create(age=30)