Added Q object support for regexes (inc. operator shortcuts)

mongoengine/queryset.py

@@ -28,6 +28,8 @@ class OperationError(Exception):
     pass
 
 
+RE_TYPE = type(re.compile(''))
+
 class Q(object):
     
     OR = '||'
@@ -47,6 +49,8 @@ class Q(object):
                 'return this.%(field)s.indexOf(a) != -1 })'),
         'size': 'this.%(field)s.length == %(value)s',
         'exists': 'this.%(field)s != null',
+        'regex_eq': '%(value)s.test(this.%(field)s)',
+        'regex_ne': '!%(value)s.test(this.%(field)s)',
     }
     
     def __init__(self, **query):
@@ -91,24 +95,36 @@ class Q(object):
                 for j, (op, value) in enumerate(value.items()):
                     # Create a custom variable name for this operator
                     op_value_name = '%so%s' % (value_name, j)
+                    # Construct the JS that uses this op
+                    operation_js = self._build_op_js(op, key, value,
+                                                     op_value_name)
                     # Update the js scope with the value for this op
                     js_scope[op_value_name] = value
-                    # Construct the JS that uses this op
-                    operation_js = Q.OPERATORS[op.strip('$')] % {
-                        'field': key, 
-                        'value': op_value_name
-                    }
                     js.append(operation_js)
             else:
-                js_scope[value_name] = value
                 # Construct the JS for this field
-                field_js = Q.OPERATORS[op.strip('$')] % {
+                field_js = self._build_op_js(op, key, value, value_name)
-                    'field': key, 
+                js_scope[value_name] = value
-                    'value': value_name
-                }
                 js.append(field_js)
         return ' && '.join(js)
 
+    def _build_op_js(self, op, key, value, value_name):
+        """Substitute the values in to the correct chunk of Javascript.
+        """
+        if isinstance(value, RE_TYPE):
+            # Regexes are handled specially
+            if op.strip('$') == 'ne':
+                op_js = Q.OPERATORS['regex_ne']
+            else:
+                op_js = Q.OPERATORS['regex_eq']
+        else:
+            op_js = Q.OPERATORS[op.strip('$')]
+        # Perform the substitution
+        operation_js = op_js % {
+            'field': key, 
+            'value': value_name
+        }
+        return operation_js
 
 class QuerySet(object):
     """A set of results returned from a query. Wraps a MongoDB cursor, 

tests/queryset.py

@@ -197,30 +197,48 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual(obj, person)
         obj = self.Person.objects(name__contains='Van').first()
         self.assertEqual(obj, None)
+        obj = self.Person.objects(Q(name__contains='van')).first()
+        self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__contains='Van')).first()
+        self.assertEqual(obj, None)
 
         # Test icontains
         obj = self.Person.objects(name__icontains='Van').first()
         self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__icontains='Van')).first()
+        self.assertEqual(obj, person)
 
         # Test startswith
         obj = self.Person.objects(name__startswith='Guido').first()
         self.assertEqual(obj, person)
         obj = self.Person.objects(name__startswith='guido').first()
         self.assertEqual(obj, None)
+        obj = self.Person.objects(Q(name__startswith='Guido')).first()
+        self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__startswith='guido')).first()
+        self.assertEqual(obj, None)
 
         # Test istartswith
         obj = self.Person.objects(name__istartswith='guido').first()
         self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__istartswith='guido')).first()
+        self.assertEqual(obj, person)
 
         # Test endswith
         obj = self.Person.objects(name__endswith='Rossum').first()
         self.assertEqual(obj, person)
         obj = self.Person.objects(name__endswith='rossuM').first()
         self.assertEqual(obj, None)
+        obj = self.Person.objects(Q(name__endswith='Rossum')).first()
+        self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__endswith='rossuM')).first()
+        self.assertEqual(obj, None)
 
         # Test iendswith
         obj = self.Person.objects(name__iendswith='rossuM').first()
         self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__iendswith='rossuM')).first()
+        self.assertEqual(obj, person)
 
     def test_filter_chaining(self):
         """Ensure filters can be chained together.
@@ -378,6 +396,26 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual(len(self.Person.objects(Q(age__in=[20]))), 2)
         self.assertEqual(len(self.Person.objects(Q(age__in=[20, 30]))), 3)
 
+    def test_q_regex(self):
+        """Ensure that Q objects can be queried using regexes.
+        """
+        person = self.Person(name='Guido van Rossum')
+        person.save()
+
+        import re
+        obj = self.Person.objects(Q(name=re.compile('^Gui'))).first()
+        self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name=re.compile('^gui'))).first()
+        self.assertEqual(obj, None)
+
+        obj = self.Person.objects(Q(name=re.compile('^gui', re.I))).first()
+        self.assertEqual(obj, person)
+
+        obj = self.Person.objects(Q(name__ne=re.compile('^bob'))).first()
+        self.assertEqual(obj, person)
+        obj = self.Person.objects(Q(name__ne=re.compile('^Gui'))).first()
+        self.assertEqual(obj, None)
+
     def test_exec_js_query(self):
         """Ensure that queries are properly formed for use in exec_js.
         """