From e2d826c412924c0baa8d446e7bd29360b3bb17a7 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Mon, 1 Oct 2012 15:26:34 +0000
Subject: [PATCH] Allow updates with match operators
 (MongoEngine/mongoengine#144)

---
 docs/changelog.rst      |  3 ++-
 mongoengine/queryset.py | 16 ++++++++++++++--
 tests/test_queryset.py  | 24 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index b749ca2b..f3e83c1b 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -4,7 +4,8 @@ Changelog
 
 Changes in 0.7.X
 ================
-- Updated URLField - can have a custom validator (MongoEngine/mongoengine#136)
+- Allow updates with match operators (MongoEngine/mongoengine#144)
+- Updated URLField - now can have a override the regex (MongoEngine/mongoengine#136)
 - Allow Django AuthenticationBackends to work with Django user (hmarr/mongoengine#573)
 - Fixed reload issue with ReferenceField where dbref=False (MongoEngine/mongoengine#138)
 
diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 5a1aa718..c774322e 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -1395,6 +1395,8 @@ class QuerySet(object):
         """
         operators = ['set', 'unset', 'inc', 'dec', 'pop', 'push', 'push_all',
                      'pull', 'pull_all', 'add_to_set']
+        match_operators = ['ne', 'gt', 'gte', 'lt', 'lte', 'in', 'nin', 'mod',
+                           'all', 'size', 'exists', 'not']
 
         mongo_update = {}
         for key, value in update.items():
@@ -1418,6 +1420,10 @@ class QuerySet(object):
                 elif op == 'add_to_set':
                     op = op.replace('_to_set', 'ToSet')
 
+            match = None
+            if parts[-1] in match_operators:
+                match = parts.pop()
+
             if _doc_cls:
                 # Switch field names to proper names [set in Field(name='foo')]
                 fields = QuerySet._lookup_field(_doc_cls, parts)
@@ -1451,16 +1457,22 @@ class QuerySet(object):
                     elif field.required or value is not None:
                         value = field.prepare_query_value(op, value)
 
+            if match:
+                match = '$' + match
+                value = {match: value}
+
             key = '.'.join(parts)
 
             if not op:
-                raise InvalidQueryError("Updates must supply an operation eg: set__FIELD=value")
+                raise InvalidQueryError("Updates must supply an operation "
+                                        "eg: set__FIELD=value")
 
             if 'pull' in op and '.' in key:
                 # Dot operators don't work on pull operations
                 # it uses nested dict syntax
                 if op == 'pullAll':
-                    raise InvalidQueryError("pullAll operations only support a single field depth")
+                    raise InvalidQueryError("pullAll operations only support "
+                                            "a single field depth")
 
                 parts.reverse()
                 for key in parts:
diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index 55531a1e..770290c9 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -414,6 +414,30 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual(post.comments[0].by, 'joe')
         self.assertEqual(post.comments[0].votes.score, 4)
 
+    def test_updates_can_have_match_operators(self):
+
+        class Post(Document):
+            title = StringField(required=True)
+            tags = ListField(StringField())
+            comments = ListField(EmbeddedDocumentField("Comment"))
+
+        class Comment(EmbeddedDocument):
+            content = StringField()
+            name = StringField(max_length=120)
+            vote = IntField()
+
+        Post.drop_collection()
+
+        comm1 = Comment(content="very funny indeed", name="John S", vote=1)
+        comm2 = Comment(content="kind of funny", name="Mark P", vote=0)
+
+        Post(title='Fun with MongoEngine', tags=['mongodb', 'mongoengine'],
+             comments=[comm1, comm2]).save()
+
+        Post.objects().update_one(pull__comments__vote__lt=1)
+
+        self.assertEqual(1, len(Post.objects.first().comments))
+
     def test_mapfield_update(self):
         """Ensure that the MapField can be updated."""
         class Member(EmbeddedDocument):