Add OR and AND operations for Q objects

src/infi/clickhouse_orm/query.py

@@ -143,9 +143,23 @@ class FOV(object):
 
 class Q(object):
 
-    def __init__(self, **kwargs):
-        self._fovs = [self._build_fov(k, v) for k, v in six.iteritems(kwargs)]
+    AND_MODE = 'AND'
+    OR_MODE = 'OR'
+
+    def __init__(self, **filter_fields):
+        self._fovs = [self._build_fov(k, v) for k, v in six.iteritems(filter_fields)]
+        self._l_child = None
+        self._r_child = None
         self._negate = False
+        self._mode = self.AND_MODE
+
+    @classmethod
+    def _construct_from(cls, l_child, r_child, mode):
+        q = Q()
+        q._l_child = l_child
+        q._r_child = r_child
+        q._mode = mode
+        return q
 
     def _build_fov(self, key, value):
         if '__' in key:
@@ -155,13 +169,23 @@ class Q(object):
         return FOV(field_name, operator, value)
 
     def to_sql(self, model_cls):
-        if not self._fovs:
-            return '1'
-        sql = ' AND '.join(fov.to_sql(model_cls) for fov in self._fovs)
+        if self._fovs:
+            sql = ' {} '.format(self._mode).join(fov.to_sql(model_cls) for fov in self._fovs)
+        else:
+            if self._l_child and self._r_child:
+                sql = '({}) {} ({})'.format(self._l_child.to_sql(model_cls), self._mode, self._r_child.to_sql(model_cls))
+            else:
+                return '1'
         if self._negate:
             sql = 'NOT (%s)' % sql
         return sql
 
+    def __or__(self, other):
+        return Q._construct_from(self, other, self.OR_MODE)
+
+    def __and__(self, other):
+        return Q._construct_from(self, other, self.AND_MODE)
+
     def __invert__(self):
         q = copy(self)
         q._negate = True
@@ -286,20 +310,24 @@ class QuerySet(object):
         qs._fields = field_names
         return qs
 
-    def filter(self, **kwargs):
+    def filter(self, *q, **filter_fields):
         """
         Returns a copy of this queryset that includes only rows matching the conditions.
+        Add q object to query if it specified.
         """
         qs = copy(self)
-        qs._q = list(self._q) + [Q(**kwargs)]
+        if q:
+            qs._q = list(self._q) + list(q)
+        else:
+            qs._q = list(self._q) + [Q(**filter_fields)]
         return qs
 
-    def exclude(self, **kwargs):
+    def exclude(self, **filter_fields):
         """
         Returns a copy of this queryset that excludes all rows matching the conditions.
         """
         qs = copy(self)
-        qs._q = list(self._q) + [~Q(**kwargs)]
+        qs._q = list(self._q) + [~Q(**filter_fields)]
         return qs
 
     def paginate(self, page_num=1, page_size=100):