Improved queryset filtering (hmarr/mongoengine#554)

mongoengine/queryset.py

@@ -4,7 +4,9 @@ import copy
 import itertools
 import operator
 
+from collections import defaultdict
 from functools import partial
+
 from mongoengine.python_support import product, reduce
 
 import pymongo
@@ -679,6 +681,7 @@ class QuerySet(object):
         custom_operators = ['match']
 
         mongo_query = {}
+        merge_query = defaultdict(list)
         for key, value in query.items():
             if key == "__raw__":
                 mongo_query.update(value)
@@ -767,21 +770,22 @@ class QuerySet(object):
             if op is None or key not in mongo_query:
                 mongo_query[key] = value
             elif key in mongo_query:
-                if isinstance(mongo_query[key], dict) and isinstance(value, dict):
+                if key in mongo_query and isinstance(mongo_query[key], dict):
                     mongo_query[key].update(value)
-                elif isinstance(mongo_query[key], list):
-                    mongo_query[key].append(value)
                 else:
-                    mongo_query[key] = [mongo_query[key], value]
+                    # Store for manually merging later
+                    merge_query[key].append(value)
 
-        for k, v in mongo_query.items():
+        # The queryset has been filter in such a way we must manually merge
+        for k, v in merge_query.items():
+            merge_query[k].append(mongo_query[k])
+            del mongo_query[k]
             if isinstance(v, list):
                 value = [{k:val} for val in v]
                 if '$and' in mongo_query.keys():
                     mongo_query['$and'].append(value)
                 else:
                     mongo_query['$and'] = value
-                del mongo_query[k]
 
         return mongo_query
 

tests/test_queryset.py

@@ -883,6 +883,21 @@ class QuerySetTest(unittest.TestCase):
         BlogPost.drop_collection()
         Blog.drop_collection()
 
+    def test_raw_and_merging(self):
+        class Doc(Document):
+            pass
+
+        raw_query = Doc.objects(__raw__={'deleted': False,
+                                'scraped': 'yes',
+                                '$nor': [{'views.extracted': 'no'},
+                                         {'attachments.views.extracted':'no'}]
+                                })._query
+
+        expected = {'deleted': False, '_types': 'Doc', 'scraped': 'yes',
+                    '$nor': [{'views.extracted': 'no'},
+                             {'attachments.views.extracted': 'no'}]}
+        self.assertEqual(expected, raw_query)
+
     def test_ordering(self):
         """Ensure default ordering is applied and can be overridden.
         """