Updated inheritable objects created by upsert now contain _cls (MongoEngine/mongoengine#118)

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 
 Changes in 0.8
 ==============
+- Updated inheritable objects created by upsert now contain _cls (MongoEngine/mongoengine#118)
 - Added support for creating documents with embedded documents in a single operation (MongoEngine/mongoengine#6)
 - Added to_json and from_json to Document (MongoEngine/mongoengine#1)
 - Added to_json and from_json to QuerySet (MongoEngine/mongoengine#131)

mongoengine/queryset/queryset.py

@@ -824,8 +824,16 @@ class QuerySet(object):
         if not write_options:
             write_options = {}
 
-        update = transform.update(self._document, **update)
         query = self._query
+        update = transform.update(self._document, **update)
+
+        # If doing an atomic upsert on an inheritable class
+        # then ensure we add _cls to the update operation
+        if upsert and '_cls' in query:
+            if '$set' in update:
+                update["$set"]["_cls"] = self._document._class_name
+            else:
+                update["$set"] = {"_cls": self._document._class_name}
 
         try:
             ret = self._collection.update(query, update, multi=multi,
@@ -852,7 +860,7 @@ class QuerySet(object):
 
         .. versionadded:: 0.2
         """
-        return self.update(safe_update=True, upsert=False, multi=False,
+        return self.update(safe_update=True, upsert=upsert, multi=False,
                            write_options=None, **update)
 
     def __iter__(self):

tests/test_queryset.py

@@ -3741,6 +3741,25 @@ class QueryFieldListTest(unittest.TestCase):
 
         self.assertEqual(doc_objects, Doc.objects.from_json(json_data))
 
+    def test_upsert_includes_cls(self):
+        """Upserts should include _cls information for inheritable classes
+        """
+
+        class Test(Document):
+            test = StringField()
+
+        Test.drop_collection()
+        Test.objects(test='foo').update_one(upsert=True, set__test='foo')
+        self.assertFalse('_cls' in Test._collection.find_one())
+
+        class Test(Document):
+            meta = {'allow_inheritance': True}
+            test = StringField()
+
+        Test.drop_collection()
+
+        Test.objects(test='foo').update_one(upsert=True, set__test='foo')
+        self.assertTrue('_cls' in Test._collection.find_one())
 
 if __name__ == '__main__':
     unittest.main()