NotUniqueError gracefully replacing ambiguous OperationError when appropriate

mongoengine/document.py

@@ -1,18 +1,19 @@
 import warnings
 
 import pymongo
+import re
 
 from bson.dbref import DBRef
 from mongoengine import signals, queryset
 
 from base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument,
                   BaseDict, BaseList)
-from queryset import OperationError
+from queryset import OperationError, NotUniqueError
 from connection import get_db, DEFAULT_CONNECTION_NAME
 
 __all__ = ['Document', 'EmbeddedDocument', 'DynamicDocument',
            'DynamicEmbeddedDocument', 'OperationError',
-           'InvalidCollectionError']
+           'InvalidCollectionError', 'NotUniqueError']
 
 
 class InvalidCollectionError(Exception):
@@ -250,8 +251,11 @@ class Document(BaseDocument):
 
         except pymongo.errors.OperationFailure, err:
             message = 'Could not save document (%s)'
-            if u'duplicate key' in unicode(err):
+            if re.match('^E1100[01] duplicate key', unicode(err)):
+                # E11000 - duplicate key error index
+                # E11001 - duplicate key on update
                 message = u'Tried to save duplicate unique keys (%s)'
+                raise NotUniqueError(message % unicode(err))
             raise OperationError(message % unicode(err))
         id_field = self._meta['id_field']
         if id_field not in self._meta.get('shard_key', []):