fix broken inheritance for Document and EmbeddedDocument

mongoengine/base/document.py

@@ -306,7 +306,7 @@ class BaseDocument(object):
             fields = []
 
         data = SON()
-        data["_id"] = None
+        data['_id'] = None
         data['_cls'] = self._class_name
 
         # only root fields ['test1.a', 'test2'] => ['test1', 'test2']
@@ -349,18 +349,8 @@ class BaseDocument(object):
                 else:
                     data[field.name] = value
 
-        # If "_id" has not been set, then try and set it
-        Document = _import_class("Document")
-        if isinstance(self, Document):
-            if data["_id"] is None:
-                data["_id"] = self._data.get("id", None)
-
-        if data['_id'] is None:
-            data.pop('_id')
-
         # Only add _cls if allow_inheritance is True
-        if (not hasattr(self, '_meta') or
+        if not self._meta.get('allow_inheritance', ALLOW_INHERITANCE):
-                not self._meta.get('allow_inheritance', ALLOW_INHERITANCE)):
             data.pop('_cls')
 
         return data

mongoengine/document.py

@@ -79,6 +79,15 @@ class EmbeddedDocument(BaseDocument):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def to_mongo(self, *args, **kwargs):
+        data = super(EmbeddedDocument, self).to_mongo(*args, **kwargs)
+
+        # remove _id from the SON if it's in it and it's None
+        if '_id' in data and data['_id'] is None:
+            del data['_id']
+
+        return data
+
     def save(self, *args, **kwargs):
         self._instance.save(*args, **kwargs)
 
@@ -204,6 +213,19 @@ class Document(BaseDocument):
                 cls.ensure_indexes()
         return cls._collection
 
+    def to_mongo(self, *args, **kwargs):
+        data = super(Document, self).to_mongo(*args, **kwargs)
+
+        # If '_id' is None, try and set it from self._data. If that
+        # doesn't exist either, remote '_id' from the SON completely.
+        if data['_id'] is None:
+            if self._data.get('id') is None:
+                del data['_id']
+            else:
+                data["_id"] = self._data['id']
+
+        return data
+
     def modify(self, query=None, **update):
         """Perform an atomic update of the document in the database and reload
         the document object using updated version.