Merge pull request #2042 from bagerard/fix_querying_embedded_subcls

Fix querying embeddedDoc sub classes

docs/changelog.rst

@@ -4,6 +4,8 @@ Changelog
 
 Development
 ===========
+- Fix querying on List(EmbeddedDocument) subclasses fields #1961 #1492
+- Fix querying on (Generic)EmbeddedDocument subclasses fields #475
 - expose `mongoengine.connection.disconnect` and `mongoengine.connection.disconnect_all`
 - Fix disconnect function #566 #1599 #605 #607 #1213 #565
 - Improve connect/disconnect documentations

mongoengine/fields.py

@@ -700,7 +700,11 @@ class EmbeddedDocumentField(BaseField):
         self.document_type.validate(value, clean)
 
     def lookup_member(self, member_name):
-        return self.document_type._fields.get(member_name)
+        doc_and_subclasses = [self.document_type] + self.document_type.__subclasses__()
+        for doc_type in doc_and_subclasses:
+            field = doc_type._fields.get(member_name)
+            if field:
+                return field
 
     def prepare_query_value(self, op, value):
         if value is not None and not isinstance(value, self.document_type):
@@ -747,12 +751,13 @@ class GenericEmbeddedDocumentField(BaseField):
         value.validate(clean=clean)
 
     def lookup_member(self, member_name):
-        if self.choices:
+        document_choices = self.choices or []
-            for choice in self.choices:
+        for document_choice in document_choices:
-                field = choice._fields.get(member_name)
+            doc_and_subclasses = [document_choice] + document_choice.__subclasses__()
+            for doc_type in doc_and_subclasses:
+                field = doc_type._fields.get(member_name)
                 if field:
                     return field
-        return None
 
     def to_mongo(self, document, use_db_field=True, fields=None):
         if document is None:

tests/fields/fields.py

@@ -8,11 +8,10 @@ from bson import DBRef, ObjectId, SON
 
 from mongoengine import Document, StringField, IntField, DateTimeField, DateField, ValidationError, \
     ComplexDateTimeField, FloatField, ListField, ReferenceField, DictField, EmbeddedDocument, EmbeddedDocumentField, \
-    GenericReferenceField, DoesNotExist, NotRegistered, GenericEmbeddedDocumentField, OperationError, DynamicField, \
+    GenericReferenceField, DoesNotExist, NotRegistered, OperationError, DynamicField, \
-    FieldDoesNotExist, EmbeddedDocumentListField, MultipleObjectsReturned, NotUniqueError, BooleanField, ObjectIdField, \
+    FieldDoesNotExist, EmbeddedDocumentListField, MultipleObjectsReturned, NotUniqueError, BooleanField,\
-    SortedListField, GenericLazyReferenceField, LazyReferenceField, DynamicDocument
+    ObjectIdField, SortedListField, GenericLazyReferenceField, LazyReferenceField, DynamicDocument
-from mongoengine.base import (BaseField, EmbeddedDocumentList,
+from mongoengine.base import (BaseField, EmbeddedDocumentList, _document_registry)
-                              _document_registry)
 
 from tests.utils import MongoDBTestCase
 
@@ -1769,79 +1768,6 @@ class FieldTest(MongoDBTestCase):
         with self.assertRaises(ValidationError):
             shirt.validate()
 
-    def test_choices_validation_documents(self):
-        """
-        Ensure fields with document choices validate given a valid choice.
-        """
-        class UserComments(EmbeddedDocument):
-            author = StringField()
-            message = StringField()
-
-        class BlogPost(Document):
-            comments = ListField(
-                GenericEmbeddedDocumentField(choices=(UserComments,))
-            )
-
-        # Ensure Validation Passes
-        BlogPost(comments=[
-            UserComments(author='user2', message='message2'),
-        ]).save()
-
-    def test_choices_validation_documents_invalid(self):
-        """
-        Ensure fields with document choices validate given an invalid choice.
-        This should throw a ValidationError exception.
-        """
-        class UserComments(EmbeddedDocument):
-            author = StringField()
-            message = StringField()
-
-        class ModeratorComments(EmbeddedDocument):
-            author = StringField()
-            message = StringField()
-
-        class BlogPost(Document):
-            comments = ListField(
-                GenericEmbeddedDocumentField(choices=(UserComments,))
-            )
-
-        # Single Entry Failure
-        post = BlogPost(comments=[
-            ModeratorComments(author='mod1', message='message1'),
-        ])
-        self.assertRaises(ValidationError, post.save)
-
-        # Mixed Entry Failure
-        post = BlogPost(comments=[
-            ModeratorComments(author='mod1', message='message1'),
-            UserComments(author='user2', message='message2'),
-        ])
-        self.assertRaises(ValidationError, post.save)
-
-    def test_choices_validation_documents_inheritance(self):
-        """
-        Ensure fields with document choices validate given subclass of choice.
-        """
-        class Comments(EmbeddedDocument):
-            meta = {
-                'abstract': True
-            }
-            author = StringField()
-            message = StringField()
-
-        class UserComments(Comments):
-            pass
-
-        class BlogPost(Document):
-            comments = ListField(
-                GenericEmbeddedDocumentField(choices=(Comments,))
-            )
-
-        # Save Valid EmbeddedDocument Type
-        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.
@@ -1958,85 +1884,6 @@ class FieldTest(MongoDBTestCase):
             self.assertEqual(error_dict['size'], SIZE_MESSAGE)
             self.assertEqual(error_dict['color'], COLOR_MESSAGE)
 
-    def test_generic_embedded_document(self):
-        class Car(EmbeddedDocument):
-            name = StringField()
-
-        class Dish(EmbeddedDocument):
-            food = StringField(required=True)
-            number = IntField()
-
-        class Person(Document):
-            name = StringField()
-            like = GenericEmbeddedDocumentField()
-
-        Person.drop_collection()
-
-        person = Person(name='Test User')
-        person.like = Car(name='Fiat')
-        person.save()
-
-        person = Person.objects.first()
-        self.assertIsInstance(person.like, Car)
-
-        person.like = Dish(food="arroz", number=15)
-        person.save()
-
-        person = Person.objects.first()
-        self.assertIsInstance(person.like, Dish)
-
-    def test_generic_embedded_document_choices(self):
-        """Ensure you can limit GenericEmbeddedDocument choices."""
-        class Car(EmbeddedDocument):
-            name = StringField()
-
-        class Dish(EmbeddedDocument):
-            food = StringField(required=True)
-            number = IntField()
-
-        class Person(Document):
-            name = StringField()
-            like = GenericEmbeddedDocumentField(choices=(Dish,))
-
-        Person.drop_collection()
-
-        person = Person(name='Test User')
-        person.like = Car(name='Fiat')
-        self.assertRaises(ValidationError, person.validate)
-
-        person.like = Dish(food="arroz", number=15)
-        person.save()
-
-        person = Person.objects.first()
-        self.assertIsInstance(person.like, Dish)
-
-    def test_generic_list_embedded_document_choices(self):
-        """Ensure you can limit GenericEmbeddedDocument choices inside
-        a list field.
-        """
-        class Car(EmbeddedDocument):
-            name = StringField()
-
-        class Dish(EmbeddedDocument):
-            food = StringField(required=True)
-            number = IntField()
-
-        class Person(Document):
-            name = StringField()
-            likes = ListField(GenericEmbeddedDocumentField(choices=(Dish,)))
-
-        Person.drop_collection()
-
-        person = Person(name='Test User')
-        person.likes = [Car(name='Fiat')]
-        self.assertRaises(ValidationError, person.validate)
-
-        person.likes = [Dish(food="arroz", number=15)]
-        person.save()
-
-        person = Person.objects.first()
-        self.assertIsInstance(person.likes[0], Dish)
-
     def test_recursive_validation(self):
         """Ensure that a validation result to_dict is available."""
         class Author(EmbeddedDocument):
@@ -2702,44 +2549,5 @@ class EmbeddedDocumentListFieldTestCase(MongoDBTestCase):
         self.assertEqual(custom_data['a'], CustomData.c_field.custom_data['a'])
 
 
-class TestEmbeddedDocumentField(MongoDBTestCase):
-    def test___init___(self):
-        class MyDoc(EmbeddedDocument):
-            name = StringField()
-
-        field = EmbeddedDocumentField(MyDoc)
-        self.assertEqual(field.document_type_obj, MyDoc)
-
-        field2 = EmbeddedDocumentField('MyDoc')
-        self.assertEqual(field2.document_type_obj, 'MyDoc')
-
-    def test___init___throw_error_if_document_type_is_not_EmbeddedDocument(self):
-        with self.assertRaises(ValidationError):
-            EmbeddedDocumentField(dict)
-
-    def test_document_type_throw_error_if_not_EmbeddedDocument_subclass(self):
-
-        class MyDoc(Document):
-            name = StringField()
-
-        emb = EmbeddedDocumentField('MyDoc')
-        with self.assertRaises(ValidationError) as ctx:
-            emb.document_type
-        self.assertIn('Invalid embedded document class provided to an EmbeddedDocumentField', str(ctx.exception))
-
-    def test_embedded_document_field_only_allow_subclasses_of_embedded_document(self):
-        # Relates to #1661
-        class MyDoc(Document):
-            name = StringField()
-
-        with self.assertRaises(ValidationError):
-            class MyFailingDoc(Document):
-                emb = EmbeddedDocumentField(MyDoc)
-
-        with self.assertRaises(ValidationError):
-            class MyFailingdoc2(Document):
-                emb = EmbeddedDocumentField('MyDoc')
-
-
 if __name__ == '__main__':
     unittest.main()

tests/fields/test_embedded_document_field.py

@@ -0,0 +1,344 @@
+# -*- coding: utf-8 -*-
+from mongoengine import Document, StringField, ValidationError, EmbeddedDocument, EmbeddedDocumentField, \
+    InvalidQueryError, LookUpError, IntField, GenericEmbeddedDocumentField, ListField, EmbeddedDocumentListField, \
+    ReferenceField
+
+from tests.utils import MongoDBTestCase
+
+
+class TestEmbeddedDocumentField(MongoDBTestCase):
+    def test___init___(self):
+        class MyDoc(EmbeddedDocument):
+            name = StringField()
+
+        field = EmbeddedDocumentField(MyDoc)
+        self.assertEqual(field.document_type_obj, MyDoc)
+
+        field2 = EmbeddedDocumentField('MyDoc')
+        self.assertEqual(field2.document_type_obj, 'MyDoc')
+
+    def test___init___throw_error_if_document_type_is_not_EmbeddedDocument(self):
+        with self.assertRaises(ValidationError):
+            EmbeddedDocumentField(dict)
+
+    def test_document_type_throw_error_if_not_EmbeddedDocument_subclass(self):
+
+        class MyDoc(Document):
+            name = StringField()
+
+        emb = EmbeddedDocumentField('MyDoc')
+        with self.assertRaises(ValidationError) as ctx:
+            emb.document_type
+        self.assertIn('Invalid embedded document class provided to an EmbeddedDocumentField', str(ctx.exception))
+
+    def test_embedded_document_field_only_allow_subclasses_of_embedded_document(self):
+        # Relates to #1661
+        class MyDoc(Document):
+            name = StringField()
+
+        with self.assertRaises(ValidationError):
+            class MyFailingDoc(Document):
+                emb = EmbeddedDocumentField(MyDoc)
+
+        with self.assertRaises(ValidationError):
+            class MyFailingdoc2(Document):
+                emb = EmbeddedDocumentField('MyDoc')
+
+    def test_query_embedded_document_attribute(self):
+        class AdminSettings(EmbeddedDocument):
+            foo1 = StringField()
+            foo2 = StringField()
+
+        class Person(Document):
+            settings = EmbeddedDocumentField(AdminSettings)
+            name = StringField()
+
+        Person.drop_collection()
+
+        p = Person(
+            settings=AdminSettings(foo1='bar1', foo2='bar2'),
+            name='John',
+        ).save()
+
+        # Test non exiting attribute
+        with self.assertRaises(InvalidQueryError) as ctx_err:
+            Person.objects(settings__notexist='bar').first()
+        self.assertEqual(unicode(ctx_err.exception), u'Cannot resolve field "notexist"')
+
+        with self.assertRaises(LookUpError):
+            Person.objects.only('settings.notexist')
+
+        # Test existing attribute
+        self.assertEqual(Person.objects(settings__foo1='bar1').first().id, p.id)
+        only_p = Person.objects.only('settings.foo1').first()
+        self.assertEqual(only_p.settings.foo1, p.settings.foo1)
+        self.assertIsNone(only_p.settings.foo2)
+        self.assertIsNone(only_p.name)
+
+        exclude_p = Person.objects.exclude('settings.foo1').first()
+        self.assertIsNone(exclude_p.settings.foo1)
+        self.assertEqual(exclude_p.settings.foo2, p.settings.foo2)
+        self.assertEqual(exclude_p.name, p.name)
+
+    def test_query_embedded_document_attribute_with_inheritance(self):
+        class BaseSettings(EmbeddedDocument):
+            meta = {'allow_inheritance': True}
+            base_foo = StringField()
+
+        class AdminSettings(BaseSettings):
+            sub_foo = StringField()
+
+        class Person(Document):
+            settings = EmbeddedDocumentField(BaseSettings)
+
+        Person.drop_collection()
+
+        p = Person(settings=AdminSettings(base_foo='basefoo', sub_foo='subfoo'))
+        p.save()
+
+        # Test non exiting attribute
+        with self.assertRaises(InvalidQueryError) as ctx_err:
+            self.assertEqual(Person.objects(settings__notexist='bar').first().id, p.id)
+        self.assertEqual(unicode(ctx_err.exception), u'Cannot resolve field "notexist"')
+
+        # Test existing attribute
+        self.assertEqual(Person.objects(settings__base_foo='basefoo').first().id, p.id)
+        self.assertEqual(Person.objects(settings__sub_foo='subfoo').first().id, p.id)
+
+        only_p = Person.objects.only('settings.base_foo', 'settings._cls').first()
+        self.assertEqual(only_p.settings.base_foo, 'basefoo')
+        self.assertIsNone(only_p.settings.sub_foo)
+
+    def test_query_list_embedded_document_with_inheritance(self):
+        class Post(EmbeddedDocument):
+            title = StringField(max_length=120, required=True)
+            meta = {'allow_inheritance': True}
+
+        class TextPost(Post):
+            content = StringField()
+
+        class MoviePost(Post):
+            author = StringField()
+
+        class Record(Document):
+            posts = ListField(EmbeddedDocumentField(Post))
+
+        record_movie = Record(posts=[MoviePost(author='John', title='foo')]).save()
+        record_text = Record(posts=[TextPost(content='a', title='foo')]).save()
+
+        records = list(Record.objects(posts__author=record_movie.posts[0].author))
+        self.assertEqual(len(records), 1)
+        self.assertEqual(records[0].id, record_movie.id)
+
+        records = list(Record.objects(posts__content=record_text.posts[0].content))
+        self.assertEqual(len(records), 1)
+        self.assertEqual(records[0].id, record_text.id)
+
+        self.assertEqual(Record.objects(posts__title='foo').count(), 2)
+
+
+class TestGenericEmbeddedDocumentField(MongoDBTestCase):
+
+    def test_generic_embedded_document(self):
+        class Car(EmbeddedDocument):
+            name = StringField()
+
+        class Dish(EmbeddedDocument):
+            food = StringField(required=True)
+            number = IntField()
+
+        class Person(Document):
+            name = StringField()
+            like = GenericEmbeddedDocumentField()
+
+        Person.drop_collection()
+
+        person = Person(name='Test User')
+        person.like = Car(name='Fiat')
+        person.save()
+
+        person = Person.objects.first()
+        self.assertIsInstance(person.like, Car)
+
+        person.like = Dish(food="arroz", number=15)
+        person.save()
+
+        person = Person.objects.first()
+        self.assertIsInstance(person.like, Dish)
+
+    def test_generic_embedded_document_choices(self):
+        """Ensure you can limit GenericEmbeddedDocument choices."""
+        class Car(EmbeddedDocument):
+            name = StringField()
+
+        class Dish(EmbeddedDocument):
+            food = StringField(required=True)
+            number = IntField()
+
+        class Person(Document):
+            name = StringField()
+            like = GenericEmbeddedDocumentField(choices=(Dish,))
+
+        Person.drop_collection()
+
+        person = Person(name='Test User')
+        person.like = Car(name='Fiat')
+        self.assertRaises(ValidationError, person.validate)
+
+        person.like = Dish(food="arroz", number=15)
+        person.save()
+
+        person = Person.objects.first()
+        self.assertIsInstance(person.like, Dish)
+
+    def test_generic_list_embedded_document_choices(self):
+        """Ensure you can limit GenericEmbeddedDocument choices inside
+        a list field.
+        """
+        class Car(EmbeddedDocument):
+            name = StringField()
+
+        class Dish(EmbeddedDocument):
+            food = StringField(required=True)
+            number = IntField()
+
+        class Person(Document):
+            name = StringField()
+            likes = ListField(GenericEmbeddedDocumentField(choices=(Dish,)))
+
+        Person.drop_collection()
+
+        person = Person(name='Test User')
+        person.likes = [Car(name='Fiat')]
+        self.assertRaises(ValidationError, person.validate)
+
+        person.likes = [Dish(food="arroz", number=15)]
+        person.save()
+
+        person = Person.objects.first()
+        self.assertIsInstance(person.likes[0], Dish)
+
+    def test_choices_validation_documents(self):
+        """
+        Ensure fields with document choices validate given a valid choice.
+        """
+        class UserComments(EmbeddedDocument):
+            author = StringField()
+            message = StringField()
+
+        class BlogPost(Document):
+            comments = ListField(
+                GenericEmbeddedDocumentField(choices=(UserComments,))
+            )
+
+        # Ensure Validation Passes
+        BlogPost(comments=[
+            UserComments(author='user2', message='message2'),
+        ]).save()
+
+    def test_choices_validation_documents_invalid(self):
+        """
+        Ensure fields with document choices validate given an invalid choice.
+        This should throw a ValidationError exception.
+        """
+        class UserComments(EmbeddedDocument):
+            author = StringField()
+            message = StringField()
+
+        class ModeratorComments(EmbeddedDocument):
+            author = StringField()
+            message = StringField()
+
+        class BlogPost(Document):
+            comments = ListField(
+                GenericEmbeddedDocumentField(choices=(UserComments,))
+            )
+
+        # Single Entry Failure
+        post = BlogPost(comments=[
+            ModeratorComments(author='mod1', message='message1'),
+        ])
+        self.assertRaises(ValidationError, post.save)
+
+        # Mixed Entry Failure
+        post = BlogPost(comments=[
+            ModeratorComments(author='mod1', message='message1'),
+            UserComments(author='user2', message='message2'),
+        ])
+        self.assertRaises(ValidationError, post.save)
+
+    def test_choices_validation_documents_inheritance(self):
+        """
+        Ensure fields with document choices validate given subclass of choice.
+        """
+        class Comments(EmbeddedDocument):
+            meta = {
+                'abstract': True
+            }
+            author = StringField()
+            message = StringField()
+
+        class UserComments(Comments):
+            pass
+
+        class BlogPost(Document):
+            comments = ListField(
+                GenericEmbeddedDocumentField(choices=(Comments,))
+            )
+
+        # Save Valid EmbeddedDocument Type
+        BlogPost(comments=[
+            UserComments(author='user2', message='message2'),
+        ]).save()
+
+    def test_query_generic_embedded_document_attribute(self):
+        class AdminSettings(EmbeddedDocument):
+            foo1 = StringField()
+
+        class NonAdminSettings(EmbeddedDocument):
+            foo2 = StringField()
+
+        class Person(Document):
+            settings = GenericEmbeddedDocumentField(choices=(AdminSettings, NonAdminSettings))
+
+        Person.drop_collection()
+
+        p1 = Person(settings=AdminSettings(foo1='bar1')).save()
+        p2 = Person(settings=NonAdminSettings(foo2='bar2')).save()
+
+        # Test non exiting attribute
+        with self.assertRaises(InvalidQueryError) as ctx_err:
+            Person.objects(settings__notexist='bar').first()
+        self.assertEqual(unicode(ctx_err.exception), u'Cannot resolve field "notexist"')
+
+        with self.assertRaises(LookUpError):
+            Person.objects.only('settings.notexist')
+
+        # Test existing attribute
+        self.assertEqual(Person.objects(settings__foo1='bar1').first().id, p1.id)
+        self.assertEqual(Person.objects(settings__foo2='bar2').first().id, p2.id)
+
+    def test_query_generic_embedded_document_attribute_with_inheritance(self):
+        class BaseSettings(EmbeddedDocument):
+            meta = {'allow_inheritance': True}
+            base_foo = StringField()
+
+        class AdminSettings(BaseSettings):
+            sub_foo = StringField()
+
+        class Person(Document):
+            settings = GenericEmbeddedDocumentField(choices=[BaseSettings])
+
+        Person.drop_collection()
+
+        p = Person(settings=AdminSettings(base_foo='basefoo', sub_foo='subfoo'))
+        p.save()
+
+        # Test non exiting attribute
+        with self.assertRaises(InvalidQueryError) as ctx_err:
+            self.assertEqual(Person.objects(settings__notexist='bar').first().id, p.id)
+        self.assertEqual(unicode(ctx_err.exception), u'Cannot resolve field "notexist"')
+
+        # Test existing attribute
+        self.assertEqual(Person.objects(settings__base_foo='basefoo').first().id, p.id)
+        self.assertEqual(Person.objects(settings__sub_foo='subfoo').first().id, p.id)