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
@@ -22,7 +23,7 @@ Changes in dev
 - Updated connection exception so it provides more info on the cause.
 - Added searching multiple levels deep in ``DictField``
 - Added ``DictField`` entries containing strings to use matching operators
-- Added ``MapField``, similar to ``DictField`` 
+- Added ``MapField``, similar to ``DictField``
 - Added Abstract Base Classes
 - Added Custom Objects Managers
 - Added sliced subfields updating
@@ -35,14 +36,14 @@ Changes in dev
 - Updated queryset to handle latest version of pymongo
   map_reduce now requires an output.
 - Added ``Document`` __hash__, __ne__ for pickling
-- Added ``FileField`` optional size arg for read method 
+- Added ``FileField`` optional size arg for read method
 - Fixed ``FileField`` seek and tell methods for reading files
-- Added ``QuerySet.clone`` to support copying querysets 
+- Added ``QuerySet.clone`` to support copying querysets
 - Fixed item_frequencies when using name thats the same as a native js function
 - Added reverse delete rules
 - Fixed issue with unset operation
 - Fixed Q-object bug
-- Added ``QuerySet.all_fields`` resets previous .only() and .exlude() 
+- Added ``QuerySet.all_fields`` resets previous .only() and .exlude()
 - Added ``QuerySet.exclude``
 - Added django style choices
 - Fixed order and filter issue
@@ -82,7 +83,7 @@ Changes in v0.3
 ===============
 - Added MapReduce support
 - Added ``contains``, ``startswith`` and ``endswith`` query operators (and
-  case-insensitive versions that are prefixed with 'i') 
+  case-insensitive versions that are prefixed with 'i')
 - Deprecated fields' ``name`` parameter, replaced with ``db_field``
 - Added ``QuerySet.only`` for only retrieving specific fields
 - Added ``QuerySet.in_bulk()`` for bulk querying using ids
@@ -129,7 +130,7 @@ Changes in v0.2
 ===============
 - Added ``Q`` class for building advanced queries
 - Added ``QuerySet`` methods for atomic updates to documents
-- Fields may now specify ``unique=True`` to enforce uniqueness across a 
+- Fields may now specify ``unique=True`` to enforce uniqueness across a
   collection
 - Added option for default document ordering
 - Fixed bug in index definitions
@@ -137,7 +138,7 @@ Changes in v0.2
 Changes in v0.1.3
 =================
 - Added Django authentication backend
-- Added ``Document.meta`` support for indexes, which are ensured just before 
+- Added ``Document.meta`` support for indexes, which are ensured just before
   querying takes place
 - A few minor bugfixes
 

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.
         """