Started work on new Q-object implementation

mongoengine/queryset.py

@@ -15,6 +15,7 @@ REPR_OUTPUT_SIZE = 20
 class DoesNotExist(Exception):
     pass
 
+
 class MultipleObjectsReturned(Exception):
     pass
 
@@ -26,12 +27,123 @@ class InvalidQueryError(Exception):
 class OperationError(Exception):
     pass
 
+
 class InvalidCollectionError(Exception):
     pass
 
+
 RE_TYPE = type(re.compile(''))
 
 
+class QNodeVisitor(object):
+
+    def visit_combination(self, combination):
+        return combination
+
+    def visit_query(self, query):
+        return query
+
+
+class SimplificationVisitor(QNodeVisitor):
+
+    def visit_combination(self, combination):
+        if combination.operation != combination.AND:
+            return combination
+
+        if any(not isinstance(node, NewQ) for node in combination.children):
+            return combination
+
+        query_ops = set()
+        query = {}
+        for node in combination.children:
+            ops = set(node.query.keys())
+            intersection = ops.intersection(query_ops)
+            if intersection:
+                msg = 'Duplicate query contitions: '
+                raise InvalidQueryError(msg + ', '.join(intersection))
+
+            query_ops.update(ops)
+            query.update(copy.deepcopy(node.query))
+        return NewQ(**query)
+
+
+class QueryCompilerVisitor(QNodeVisitor):
+
+    def __init__(self, document):
+        self.document = document
+
+    def visit_combination(self, combination):
+        if combination.operation == combination.OR:
+            return combination
+        return combination
+
+    def visit_query(self, query):
+        return QuerySet._transform_query(self.document, **query.query)
+
+
+class QNode(object):
+
+    AND = 0
+    OR = 1
+
+    def to_query(self, document):
+        query = self.accept(SimplificationVisitor())
+        query = query.accept(QueryCompilerVisitor(document))
+        return query
+
+    def accept(self, visitor):
+        raise NotImplementedError
+
+    def _combine(self, other, operation):
+        if other.empty:
+            return self
+
+        if self.empty:
+            return other
+
+        return QCombination(operation, [self, other])
+
+    @property
+    def empty(self):
+        return False
+
+    def __or__(self, other):
+        return self._combine(other, self.OR)
+
+    def __and__(self, other):
+        return self._combine(other, self.AND)
+
+
+class QCombination(QNode):
+
+    def __init__(self, operation, children):
+        self.operation = operation
+        self.children = children
+
+    def accept(self, visitor):
+        for i in range(len(self.children)):
+            self.children[i] = self.children[i].accept(visitor)
+
+        return visitor.visit_combination(self)
+
+    @property
+    def empty(self):
+        return not bool(self.query)
+
+
+class NewQ(QNode):
+
+    def __init__(self, **query):
+        self.query = query
+
+    def accept(self, visitor):
+        return visitor.visit_query(self)
+
+    @property
+    def empty(self):
+        return not bool(self.query)
+
+
 class Q(object):
 
     OR = '||'

tests/queryset.py

@@ -6,7 +6,7 @@ import pymongo
 from datetime import datetime, timedelta
 
 from mongoengine.queryset import (QuerySet, MultipleObjectsReturned,
-                                  DoesNotExist)
+                                  DoesNotExist, NewQ)
 from mongoengine import *
 
 
@@ -53,9 +53,6 @@ class QuerySetTest(unittest.TestCase):
         person2 = self.Person(name="User B", age=30)
         person2.save()
 
-        q1 = Q(name='test')
-        q2 = Q(age__gte=18)
-
         # Find all people in the collection
         people = self.Person.objects
         self.assertEqual(len(people), 2)
@@ -1415,5 +1412,25 @@ class QTest(unittest.TestCase):
         self.assertEqual(Post.objects.filter(Q(created_user=user)).count(), 1)
 
 
+class NewQTest(unittest.TestCase):
+
+    def test_and_combination(self):
+        class TestDoc(Document):
+            x = IntField()
+
+        # Check than an error is raised when conflicting queries are anded
+        def invalid_combination():
+            query = NewQ(x__lt=7) & NewQ(x__lt=3)
+            query.to_query(TestDoc)
+        self.assertRaises(InvalidQueryError, invalid_combination)
+
+        # Check normal cases work without an error
+        query = NewQ(x__lt=7) & NewQ(x__gt=3)
+
+        q1 = NewQ(x__lt=7)
+        q2 = NewQ(x__gt=3)
+        query = (q1 & q2).to_query(TestDoc)
+        self.assertEqual(query, {'x': {'$lt': 7, '$gt': 3}})
+
 if __name__ == '__main__':
     unittest.main()