Inconsistent setting of '_cls' broke inherited document referencing

Fixes #199

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 Changes in dev
 ==============
 
+- Fixed issue with inconsitent setting of _cls breaking inherited referencing
 - Added help_text and verbose_name to fields to help with some form libs
 - Updated item_frequencies to handle embedded document lookups
 - Added delta tracking now only sets / unsets explicitly changed fields

mongoengine/base.py

@@ -173,7 +173,7 @@ class ComplexBaseField(BaseField):
 
         for k,v in value_list.items():
             if isinstance(v, dict) and '_cls' in v and '_ref' not in v:
-                value_list[k] = get_document(v['_cls'].split('.')[-1])._from_son(v)
+                value_list[k] = get_document(v['_cls'])._from_son(v)
 
         # Handle all dereferencing
         db = _get_db()
@@ -401,6 +401,7 @@ class DocumentMetaclass(type):
                 else:
                     simple_class = False
 
+        doc_class_name = '.'.join(reversed(class_name))
         meta = attrs.get('_meta', attrs.get('meta', {}))
 
         if 'allow_inheritance' not in meta:
@@ -412,8 +413,7 @@ class DocumentMetaclass(type):
             raise ValueError('Only direct subclasses of Document may set '
                              '"allow_inheritance" to False')
         attrs['_meta'] = meta
-
-        attrs['_class_name'] = '.'.join(reversed(class_name))
+        attrs['_class_name'] = doc_class_name
         attrs['_superclasses'] = superclasses
 
         # Add the document's fields to the _fields attribute
@@ -448,7 +448,7 @@ class DocumentMetaclass(type):
         new_class.add_to_class('MultipleObjectsReturned', exc)
 
         global _document_registry
-        _document_registry[name] = new_class
+        _document_registry[doc_class_name] = new_class
 
         return new_class
 

mongoengine/fields.py

@@ -652,7 +652,7 @@ class GenericReferenceField(BaseField):
         id_ = id_field.to_mongo(id_)
         collection = document._meta['collection']
         ref = pymongo.dbref.DBRef(collection, id_)
-        return {'_cls': document.__class__.__name__, '_ref': ref}
+        return {'_cls': document._class_name, '_ref': ref}
 
     def prepare_query_value(self, op, value):
         return self.to_mongo(value)

tests/document.py

@@ -116,6 +116,8 @@ class DocumentTest(unittest.TestCase):
         class Human(Mammal): pass
         class Dog(Mammal): pass
 
+        Animal.drop_collection()
+
         Animal().save()
         Fish().save()
         Mammal().save()
@@ -133,6 +135,52 @@ class DocumentTest(unittest.TestCase):
 
         Animal.drop_collection()
 
+    def test_polymorphic_references(self):
+        """Ensure that the correct subclasses are returned from a query when
+        using references / generic references
+        """
+        class Animal(Document): pass
+        class Fish(Animal): pass
+        class Mammal(Animal): pass
+        class Human(Mammal): pass
+        class Dog(Mammal): pass
+
+        class Zoo(Document):
+            animals = ListField(ReferenceField(Animal))
+
+        Zoo.drop_collection()
+        Animal.drop_collection()
+
+        Animal().save()
+        Fish().save()
+        Mammal().save()
+        Human().save()
+        Dog().save()
+
+        # Save a reference to each animal
+        zoo = Zoo(animals=Animal.objects)
+        zoo.save()
+        zoo.reload()
+
+        classes = [a.__class__ for a in Zoo.objects.first().animals]
+        self.assertEqual(classes, [Animal, Fish, Mammal, Human, Dog])
+
+        Zoo.drop_collection()
+
+        class Zoo(Document):
+            animals = ListField(GenericReferenceField(Animal))
+
+        # Save a reference to each animal
+        zoo = Zoo(animals=Animal.objects)
+        zoo.save()
+        zoo.reload()
+
+        classes = [a.__class__ for a in Zoo.objects.first().animals]
+        self.assertEqual(classes, [Animal, Fish, Mammal, Human, Dog])
+
+        Zoo.drop_collection()
+        Animal.drop_collection()
+
     def test_inheritance(self):
         """Ensure that document may inherit fields from a superclass document.
         """