Fix GenericReferenceField choices with DBRef and let it possible to set Document choice as string

mongoengine/base/fields.py

@@ -165,16 +165,13 @@ class BaseField(object):
         """
         pass
 
-    def _validate(self, value, **kwargs):
+    def _validate_choices(self, value):
         Document = _import_class('Document')
         EmbeddedDocument = _import_class('EmbeddedDocument')
 
-        # Check the Choices Constraint
-        if self.choices:
-
         choice_list = self.choices
-            if isinstance(self.choices[0], (list, tuple)):
-                choice_list = [k for k, v in self.choices]
+        if isinstance(choice_list[0], (list, tuple)):
+            choice_list = [k for k, v in choice_list]
 
         # Choices which are other types of Documents
         if isinstance(value, (Document, EmbeddedDocument)):
@@ -186,6 +183,12 @@ class BaseField(object):
         elif value not in choice_list:
             self.error('Value must be one of %s' % unicode(choice_list))
 
+
+    def _validate(self, value, **kwargs):
+        # Check the Choices Constraint
+        if self.choices:
+            self._validate_choices(value)
+
         # check validation argument
         if self.validation is not None:
             if callable(self.validation):

mongoengine/fields.py

@@ -1140,6 +1140,35 @@ class GenericReferenceField(BaseField):
     .. versionadded:: 0.3
     """
 
+    def __init__(self, *args, **kwargs):
+        choices = kwargs.pop('choices', None)
+        super(GenericReferenceField, self).__init__(*args, **kwargs)
+        self._original_choices = choices or []
+        self._cooked_choices = None
+
+    def _validate_choices(self, value):
+        if isinstance(value, dict):
+            # If the field has not been dereferenced, it is still a dict
+            # of class and DBRef
+            if value.get('_cls') in [c.__name__ for c in self.choices]:
+                return
+        super(GenericReferenceField, self)._validate_choices(value)
+
+    @property
+    def choices(self):
+        if self._cooked_choices is None:
+            self._cooked_choices = []
+            for choice in self._original_choices:
+                if isinstance(choice, basestring):
+                    choice = get_document(choice)
+                self._cooked_choices.append(choice)
+        return self._cooked_choices
+
+    @choices.setter
+    def choices(self, value):
+        self._original_choices = value
+        self._cooked_choices = None
+
     def __get__(self, instance, owner):
         if instance is None:
             return self

tests/fields/fields.py

@@ -2396,6 +2396,62 @@ class FieldTest(unittest.TestCase):
         bm = Bookmark.objects.first()
         self.assertEqual(bm.bookmark_object, post_1)
 
+    def test_generic_reference_string_choices(self):
+        """Ensure that a GenericReferenceField can handle choices as strings
+        """
+        class Link(Document):
+            title = StringField()
+
+        class Post(Document):
+            title = StringField()
+
+        class Bookmark(Document):
+            bookmark_object = GenericReferenceField(choices=('Post', Link))
+
+        Link.drop_collection()
+        Post.drop_collection()
+        Bookmark.drop_collection()
+
+        link_1 = Link(title="Pitchfork")
+        link_1.save()
+
+        post_1 = Post(title="Behind the Scenes of the Pavement Reunion")
+        post_1.save()
+
+        bm = Bookmark(bookmark_object=link_1)
+        bm.save()
+
+        bm = Bookmark(bookmark_object=post_1)
+        bm.save()
+
+        bm = Bookmark(bookmark_object=bm)
+        self.assertRaises(ValidationError, bm.validate)
+
+    def test_generic_reference_choices_no_dereference(self):
+        """Ensure that a GenericReferenceField can handle choices on
+        non-derefenreced (i.e. DBRef) elements
+        """
+        class Post(Document):
+            title = StringField()
+
+        class Bookmark(Document):
+            bookmark_object = GenericReferenceField(choices=(Post, ))
+            other_field = StringField()
+
+        Post.drop_collection()
+        Bookmark.drop_collection()
+
+        post_1 = Post(title="Behind the Scenes of the Pavement Reunion")
+        post_1.save()
+
+        bm = Bookmark(bookmark_object=post_1)
+        bm.save()
+
+        bm = Bookmark.objects.get(id=bm.id)
+        # bookmark_object is now a DBRef
+        bm.other_field = 'dummy_change'
+        bm.save()
+
     def test_generic_reference_list_choices(self):
         """Ensure that a ListField properly dereferences generic references and
         respects choices.