DoesNotExist and MultipleObjectsReturned now contributes Document class

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})

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.

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)