#2811

add array filter to modify

Author: idocyabra

mongoengine/queryset/base.py

@@ -634,7 +634,7 @@ def upsert_one(self, write_concern=None, read_concern=None, **update):
             document = self._document.objects.with_id(atomic_update.upserted_id)
         return document
 
-    def update_one(self, upsert=False, write_concern=None, full_result=False, **update):
+    def update_one(self, upsert=False, write_concern=None, full_result=False, array_filters=None, **update):
         """Perform an atomic update on the fields of the first document
         matched by the query.
 
@@ -647,6 +647,7 @@ def update_one(self, upsert=False, write_concern=None, full_result=False, **upda
             will force an fsync on the primary server.
         :param full_result: Return the associated ``pymongo.UpdateResult`` rather than just the number
             updated items
+        :param array_filters: A list of filters specifying which array elements an update should apply.
         :param update: Django-style update keyword arguments
             full_result
         :returns the number of updated documents (unless ``full_result`` is True)
@@ -656,11 +657,14 @@ def update_one(self, upsert=False, write_concern=None, full_result=False, **upda
             multi=False,
             write_concern=write_concern,
             full_result=full_result,
+            array_filters=array_filters,
             **update,
         )
 
     def modify(
-        self, upsert=False, full_response=False, remove=False, new=False, **update
+        self, upsert=False, full_response=False, remove=False, new=False,
+            array_filters=None,
+            **update
     ):
         """Update and return the updated document.
 
@@ -680,6 +684,7 @@ def modify(
         :param remove: remove rather than updating (default ``False``)
         :param new: return updated rather than original document
             (default ``False``)
+        :param array_filters: A list of filters specifying which array elements an update should apply.
         :param update: Django-style update keyword arguments
         """
 
@@ -717,6 +722,7 @@ def modify(
                     upsert=upsert,
                     sort=sort,
                     return_document=return_doc,
+                    array_filters=array_filters,
                     **self._cursor_args,
                 )
         except pymongo.errors.DuplicateKeyError as err:

tests/queryset/test_queryset.py

@@ -615,6 +615,22 @@ class Blog(Document):
 
         assert testc_blogs.count() == 1
 
+        # modify
+        Blog.drop_collection()
+
+        # update one
+        Blog.objects.create(tags=["test1", "test2", "test3"])
+
+        new_blog = Blog.objects().modify(
+            __raw__={"$set": {"tags.$[element]": "test11111"}},
+            array_filters=[{"element": {"$eq": "test2"}}],
+            new=True,
+        )
+        testc_blogs = Blog.objects(tags="test11111")
+        assert new_blog == testc_blogs.first()
+
+        assert testc_blogs.count() == 1
+
         Blog.drop_collection()
 
         # update one inner list