Compilation of combinations - simple $or now works

This commit is contained in:
Harry Marr 2010-10-03 21:26:26 +01:00
parent 62388cb740
commit a3c46fec07
2 changed files with 33 additions and 13 deletions

View File

@ -43,6 +43,20 @@ class QNodeVisitor(object):
def visit_query(self, query):
return query
def _query_conjunction(self, queries):
query_ops = set()
combined_query = {}
for query in queries:
ops = set(query.keys())
intersection = ops.intersection(query_ops)
if intersection:
msg = 'Duplicate query contitions: '
raise InvalidQueryError(msg + ', '.join(intersection))
query_ops.update(ops)
combined_query.update(copy.deepcopy(query))
return combined_query
class SimplificationVisitor(QNodeVisitor):
@ -53,18 +67,8 @@ class SimplificationVisitor(QNodeVisitor):
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)
queries = [node.query for node in combination.children]
return NewQ(**self._query_conjunction(queries))
class QueryCompilerVisitor(QNodeVisitor):
@ -74,7 +78,9 @@ class QueryCompilerVisitor(QNodeVisitor):
def visit_combination(self, combination):
if combination.operation == combination.OR:
return combination
return {'$or': combination.children}
elif combination.operation == combination.AND:
return self._query_conjunction(combination.children)
return combination
def visit_query(self, query):

View File

@ -1432,5 +1432,19 @@ class NewQTest(unittest.TestCase):
query = (q1 & q2).to_query(TestDoc)
self.assertEqual(query, {'x': {'$lt': 7, '$gt': 3}})
def test_or_combination(self):
class TestDoc(Document):
x = IntField()
q1 = NewQ(x__lt=3)
q2 = NewQ(x__gt=7)
query = (q1 | q2).to_query(TestDoc)
self.assertEqual(query, {
'$or': [
{'x': {'$lt': 3}},
{'x': {'$gt': 7}},
]
})
if __name__ == '__main__':
unittest.main()