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.