From 1a24d599b31e020733f6a9e8a64a79579e6bd8d9 Mon Sep 17 00:00:00 2001
From: Matthew Ellison <seglberg@gmail.com>
Date: Sun, 11 Jan 2015 19:12:10 -0500
Subject: [PATCH] Field Choices Now Accept Subclasses of Documents

- Fields containing 'choices' of which a choice is an
  EmbeddedDocument or Document will now accept subclasses of that
  choice.
---
 docs/changelog.rst         |  1 +
 mongoengine/base/fields.py | 28 +++++++++++++------------
 tests/fields/fields.py     | 43 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index fc8b281c..ae4ba85a 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -5,6 +5,7 @@ Changelog
 
 Changes in 0.9.X - DEV
 ======================
+- Field Choices Now Accept Subclasses of Documents
 - Ensure Indexes before Each Save #812
 - Generate Unique Indices for Lists of EmbeddedDocuments #358
 - Sparse fields #515
diff --git a/mongoengine/base/fields.py b/mongoengine/base/fields.py
index 6fad5d63..359ea6d2 100644
--- a/mongoengine/base/fields.py
+++ b/mongoengine/base/fields.py
@@ -158,21 +158,23 @@ class BaseField(object):
     def _validate(self, value, **kwargs):
         Document = _import_class('Document')
         EmbeddedDocument = _import_class('EmbeddedDocument')
-        # check choices
+
+        # Check the Choices Constraint
         if self.choices:
-            is_cls = isinstance(value, (Document, EmbeddedDocument))
-            value_to_check = value.__class__ if is_cls else value
-            err_msg = 'an instance' if is_cls else 'one'
+
+            choice_list = self.choices
             if isinstance(self.choices[0], (list, tuple)):
-                option_keys = [k for k, v in self.choices]
-                if value_to_check not in option_keys:
-                    msg = ('Value must be %s of %s' %
-                           (err_msg, unicode(option_keys)))
-                    self.error(msg)
-            elif value_to_check not in self.choices:
-                msg = ('Value must be %s of %s' %
-                       (err_msg, unicode(self.choices)))
-                self.error(msg)
+                choice_list = [k for k, v in self.choices]
+
+            # Choices which are other types of Documents
+            if isinstance(value, (Document, EmbeddedDocument)):
+                if not any(isinstance(value, c) for c in choice_list):
+                    self.error(
+                        'Value must be instance of %s' % unicode(choice_list)
+                    )
+            # Choices which are types other than Documents
+            elif value not in choice_list:
+                self.error('Value must be one of %s' % unicode(choice_list))
 
         # check validation argument
         if self.validation is not None:
diff --git a/tests/fields/fields.py b/tests/fields/fields.py
index a95001b4..fedc7c38 100644
--- a/tests/fields/fields.py
+++ b/tests/fields/fields.py
@@ -2446,6 +2446,49 @@ class FieldTest(unittest.TestCase):
 
         Shirt.drop_collection()
 
+    def test_choices_validation_documents(self):
+        """
+        Ensure that a field with choices which are documents are
+        validated correctly.
+        """
+        class UserComments(EmbeddedDocument):
+            author = StringField()
+            message = StringField()
+
+        class BlogPost(Document):
+            comments = ListField(
+                GenericEmbeddedDocumentField(choices=(UserComments,))
+            )
+
+        BlogPost(comments=[
+            UserComments(author='user2', message='message2'),
+        ]).save()
+
+    def test_choices_validation_documents_inheritance(self):
+        """
+        Ensure that a field with choices which are documents are
+        validated correctly when a subclass of a valid choice is used.
+        """
+        class Comments(EmbeddedDocument):
+            meta = {
+                'abstract': True
+            }
+
+            author = StringField()
+            message = StringField()
+
+        class UserComments(Comments):
+            pass
+
+        class BlogPost(Document):
+            comments = ListField(
+                GenericEmbeddedDocumentField(choices=(Comments,))
+            )
+
+        BlogPost(comments=[
+            UserComments(author='user2', message='message2'),
+        ]).save()
+
     def test_choices_get_field_display(self):
         """Test dynamic helper for returning the display value of a choices
         field.