fix-#397: Allow specifying the '_cls' as a field for indexes

mongoengine/base/document.py

@@ -26,7 +26,7 @@ NON_FIELD_ERRORS = '__all__'
 
 class BaseDocument(object):
     __slots__ = ('_changed_fields', '_initialised', '_created', '_data',
-                 '_dynamic_fields', '_auto_id_field', '_db_field_map', '_cls', '__weakref__')
+                 '_dynamic_fields', '_auto_id_field', '_db_field_map', '__weakref__')
 
     _dynamic = False
     _dynamic_lock = True
@@ -78,6 +78,9 @@ class BaseDocument(object):
             value = getattr(self, key, None)
             setattr(self, key, value)
 
+        if "_cls" not in values:
+            self._cls = self._class_name
+
         # Set passed values after initialisation
         if self._dynamic:
             dynamic_data = {}

mongoengine/base/metaclasses.py

@@ -47,6 +47,10 @@ class DocumentMetaclass(type):
                     meta.merge(base._meta)
             attrs['_meta'] = meta
 
+        if '_meta' in attrs and attrs['_meta'].get('allow_inheritance', ALLOW_INHERITANCE):
+            StringField = _import_class('StringField')
+            attrs['_cls'] = StringField()
+
         # Handle document Fields
 
         # Merge all fields from subclasses

mongoengine/dereference.py

@@ -180,7 +180,11 @@ class DeReference(object):
                 return self.object_map.get(items['_ref'].id, items)
             elif '_cls' in items:
                 doc = get_document(items['_cls'])._from_son(items)
+                _cls = doc._data.pop('_cls', None)
+                del items['_cls']
                 doc._data = self._attach_objects(doc._data, depth, doc, None)
+                if _cls is not None:
+                    doc._data['_cls'] = _cls
                 return doc
 
         if not hasattr(items, 'items'):

mongoengine/queryset/base.py

@@ -1413,8 +1413,11 @@ class BaseQuerySet(object):
     def _query(self):
         if self._mongo_query is None:
             self._mongo_query = self._query_obj.to_query(self._document)
-            if self._class_check:
-                self._mongo_query.update(self._initial_query)
+            if self._class_check and self._initial_query:
+                if "_cls" in self._mongo_query:
+                    self._mongo_query = {"$and": [self._initial_query, self._mongo_query]}
+                else:
+                    self._mongo_query.update(self._initial_query)
         return self._mongo_query
 
     @property