diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 3583d0f9..087ee487 100644 --- a/mongoengine/queryset.py +++ b/mongoengine/queryset.py @@ -522,6 +522,15 @@ class QuerySet(object): fields = [] field = None for field_name in parts: + # Handle ListField indexing: + if field_name.isdigit(): + try: + field = field.field + except AttributeError, err: + raise InvalidQueryError( + "Can't use index on unsubscriptable field (%s)" % err) + fields.append(field_name) + continue if field is None: # Look up first field from the document if field_name == 'pk': @@ -1070,10 +1079,17 @@ class QuerySet(object): if _doc_cls: # Switch field names to proper names [set in Field(name='foo')] fields = QuerySet._lookup_field(_doc_cls, parts) - parts = [field.db_field for field in fields] + parts = [] + + for field in fields: + if isinstance(field, str): + parts.append(field) + else: + parts.append(field.db_field) # Convert value to proper value field = fields[-1] + if op in (None, 'set', 'push', 'pull', 'addToSet'): value = field.prepare_query_value(op, value) elif op in ('pushAll', 'pullAll'): diff --git a/tests/queryset.py b/tests/queryset.py index 6c0b686f..091a994b 100644 --- a/tests/queryset.py +++ b/tests/queryset.py @@ -212,6 +212,55 @@ class QuerySetTest(unittest.TestCase): Blog.drop_collection() + def test_update_array_position(self): + """Ensure that updating by array position works. + + Check update() and update_one() can take syntax like: + set__posts__1__comments__1__name="testc" + Check that it only works for ListFields. + """ + class Comment(EmbeddedDocument): + name = StringField() + + class Post(EmbeddedDocument): + comments = ListField(EmbeddedDocumentField(Comment)) + + class Blog(Document): + tags = ListField(StringField()) + posts = ListField(EmbeddedDocumentField(Post)) + + Blog.drop_collection() + + comment1 = Comment(name='testa') + comment2 = Comment(name='testb') + post1 = Post(comments=[comment1, comment2]) + post2 = Post(comments=[comment2, comment2]) + blog1 = Blog.objects.create(posts=[post1, post2]) + blog2 = Blog.objects.create(posts=[post2, post1]) + + # Update all of the first comments of second posts of all blogs + blog = Blog.objects().update(set__posts__1__comments__0__name="testc") + testc_blogs = Blog.objects(posts__1__comments__0__name="testc") + self.assertEqual(len(testc_blogs), 2) + + Blog.drop_collection() + + blog1 = Blog.objects.create(posts=[post1, post2]) + blog2 = Blog.objects.create(posts=[post2, post1]) + + # Update only the first blog returned by the query + blog = Blog.objects().update_one( + set__posts__1__comments__1__name="testc") + testc_blogs = Blog.objects(posts__1__comments__1__name="testc") + self.assertEqual(len(testc_blogs), 1) + + # Check that using this indexing syntax on a non-list fails + def non_list_indexing(): + Blog.objects().update(set__posts__1__comments__0__name__1="asdf") + self.assertRaises(InvalidQueryError, non_list_indexing) + + Blog.drop_collection() + def test_get_or_create(self): """Ensure that ``get_or_create`` returns one result or creates a new document.