Merge pull request #1655 from werat/fix_update_pull_in

Fix pull+in update queries

AUTHORS

@@ -244,4 +244,5 @@ that much better:
  * Stanislav Kaledin (https://github.com/sallyruthstruik)
  * Dmitry Yantsen (https://github.com/mrTable)
  * Renjianxin (https://github.com/Davidrjx)
- * Erdenezul Batmunkh (https://github.com/erdenezul)
+ * Erdenezul Batmunkh (https://github.com/erdenezul)
+ * Andy Yankovsky (https://github.com/werat)

mongoengine/queryset/transform.py

@@ -101,21 +101,8 @@ def query(_doc_cls=None, **kwargs):
                         value = value['_id']
 
             elif op in ('in', 'nin', 'all', 'near') and not isinstance(value, dict):
-                # Raise an error if the in/nin/all/near param is not iterable. We need a
+                # Raise an error if the in/nin/all/near param is not iterable.
-                # special check for BaseDocument, because - although it's iterable - using
+                value = _prepare_query_for_iterable(field, op, value)
-                # it as such in the context of this method is most definitely a mistake.
-                BaseDocument = _import_class('BaseDocument')
-                if isinstance(value, BaseDocument):
-                    raise TypeError("When using the `in`, `nin`, `all`, or "
-                                    "`near`-operators you can\'t use a "
-                                    "`Document`, you must wrap your object "
-                                    "in a list (object -> [object]).")
-                elif not hasattr(value, '__iter__'):
-                    raise TypeError("The `in`, `nin`, `all`, or "
-                                    "`near`-operators must be applied to an "
-                                    "iterable (e.g. a list).")
-                else:
-                    value = [field.prepare_query_value(op, v) for v in value]
 
             # If we're querying a GenericReferenceField, we need to alter the
             # key depending on the value:
@@ -284,9 +271,15 @@ def update(_doc_cls=None, **update):
             if isinstance(field, GeoJsonBaseField):
                 value = field.to_mongo(value)
 
-            if op == 'push' and isinstance(value, (list, tuple, set)):
+            if op == 'pull':
+                if field.required or value is not None:
+                    if match == 'in' and not isinstance(value, dict):
+                        value = _prepare_query_for_iterable(field, op, value)
+                    else:
+                        value = field.prepare_query_value(op, value)
+            elif op == 'push' and isinstance(value, (list, tuple, set)):
                 value = [field.prepare_query_value(op, v) for v in value]
-            elif op in (None, 'set', 'push', 'pull'):
+            elif op in (None, 'set', 'push'):
                 if field.required or value is not None:
                     value = field.prepare_query_value(op, value)
             elif op in ('pushAll', 'pullAll'):
@@ -437,3 +430,22 @@ def _infer_geometry(value):
 
     raise InvalidQueryError('Invalid $geometry data. Can be either a '
                             'dictionary or (nested) lists of coordinate(s)')
+
+
+def _prepare_query_for_iterable(field, op, value):
+    # We need a special check for BaseDocument, because - although it's iterable - using
+    # it as such in the context of this method is most definitely a mistake.
+    BaseDocument = _import_class('BaseDocument')
+
+    if isinstance(value, BaseDocument):
+        raise TypeError("When using the `in`, `nin`, `all`, or "
+                        "`near`-operators you can\'t use a "
+                        "`Document`, you must wrap your object "
+                        "in a list (object -> [object]).")
+
+    if not hasattr(value, '__iter__'):
+        raise TypeError("The `in`, `nin`, `all`, or "
+                        "`near`-operators must be applied to an "
+                        "iterable (e.g. a list).")
+
+    return [field.prepare_query_value(op, v) for v in value]

tests/document/instance.py

@@ -1341,6 +1341,23 @@ class InstanceTest(unittest.TestCase):
         site = Site.objects.first()
         self.assertEqual(site.page.log_message, "Error: Dummy message")
 
+    def test_update_list_field(self):
+        """Test update on `ListField` with $pull + $in.
+        """
+        class Doc(Document):
+            foo = ListField(StringField())
+
+        Doc.drop_collection()
+        doc = Doc(foo=['a', 'b', 'c'])
+        doc.save()
+
+        # Update
+        doc = Doc.objects.first()
+        doc.update(pull__foo__in=['a', 'c'])
+
+        doc = Doc.objects.first()
+        self.assertEqual(doc.foo, ['b'])
+
     def test_embedded_update_db_field(self):
         """Test update on `EmbeddedDocumentField` fields when db_field
         is other than default.

tests/queryset/transform.py

@@ -28,12 +28,16 @@ class TransformTest(unittest.TestCase):
                          {'name': {'$exists': True}})
 
     def test_transform_update(self):
+        class LisDoc(Document):
+            foo = ListField(StringField())
+
         class DicDoc(Document):
             dictField = DictField()
 
         class Doc(Document):
             pass
 
+        LisDoc.drop_collection()
         DicDoc.drop_collection()
         Doc.drop_collection()
 
@@ -51,6 +55,9 @@ class TransformTest(unittest.TestCase):
         update = transform.update(DicDoc, pull__dictField__test=doc)
         self.assertTrue(isinstance(update["$pull"]["dictField"]["test"], dict))
 
+        update = transform.update(LisDoc, pull__foo__in=['a'])
+        self.assertEqual(update, {'$pull': {'foo': {'$in': ['a']}}})
+
     def test_query_field_name(self):
         """Ensure that the correct field name is used when querying.
         """