Refactor: move WhereNode.as_mql_idx to index.py

Author: WaVEV

django_mongodb/indexes.py

@@ -1,5 +1,8 @@
+from django.core.exceptions import EmptyResultSet, FullResultSet
+from django.db import NotSupportedError
 from django.db.models import Index
 from django.db.models.sql.query import Query
+from django.db.models.sql.where import AND, XOR, WhereNode
 
 
 def _get_condition_mql(self, model, schema_editor):
@@ -10,5 +13,49 @@ def _get_condition_mql(self, model, schema_editor):
     return where.as_mql_idx(compiler, schema_editor.connection)
 
 
+def where_node_idx(self, compiler, connection):
+    if self.connector == AND:
+        full_needed, empty_needed = len(self.children), 1
+    else:
+        full_needed, empty_needed = 1, len(self.children)
+    if self.connector == AND:
+        operator = "$and"
+    elif self.connector == XOR:
+        raise NotSupportedError("Xor in indexes is not supported")
+    else:
+        operator = "$or"
+    if self.negated:
+        raise NotSupportedError("Negated field in indexes is not supported")
+    children_mql = []
+    for child in self.children:
+        try:
+            mql = child.as_mql_idx(compiler, connection)
+        except EmptyResultSet:
+            empty_needed -= 1
+        except FullResultSet:
+            full_needed -= 1
+        else:
+            if mql:
+                children_mql.append(mql)
+            else:
+                full_needed -= 1
+
+    if empty_needed == 0:
+        raise EmptyResultSet
+    if full_needed == 0:
+        return {}
+
+    if len(children_mql) == 1:
+        mql = children_mql[0]
+    elif len(children_mql) > 1:
+        mql = {operator: children_mql}
+    else:
+        mql = {}
+    if not mql:
+        raise FullResultSet
+    return mql
+
+
 def register_indexes():
     Index._get_condition_mql = _get_condition_mql
+    WhereNode.as_mql_idx = where_node_idx

django_mongodb/query.py

@@ -297,52 +297,8 @@ def where_node(self, compiler, connection):
     return mql
 
 
-def where_node_idx(self, compiler, connection):
-    if self.connector == AND:
-        full_needed, empty_needed = len(self.children), 1
-    else:
-        full_needed, empty_needed = 1, len(self.children)
-    if self.connector == AND:
-        operator = "$and"
-    elif self.connector == XOR:
-        raise NotSupportedError("Xor in indexes is not supported")
-    else:
-        operator = "$or"
-    if self.negated:
-        raise NotSupportedError("Negated field in indexes is not supported")
-    children_mql = []
-    for child in self.children:
-        try:
-            mql = child.as_mql_idx(compiler, connection)
-        except EmptyResultSet:
-            empty_needed -= 1
-        except FullResultSet:
-            full_needed -= 1
-        else:
-            if mql:
-                children_mql.append(mql)
-            else:
-                full_needed -= 1
-
-    if empty_needed == 0:
-        raise EmptyResultSet
-    if full_needed == 0:
-        return {}
-
-    if len(children_mql) == 1:
-        mql = children_mql[0]
-    elif len(children_mql) > 1:
-        mql = {operator: children_mql}
-    else:
-        mql = {}
-    if not mql:
-        raise FullResultSet
-    return mql
-
-
 def register_nodes():
     ExtraWhere.as_mql = extra_where
     Join.as_mql = join
     NothingNode.as_mql = NothingNode.as_sql
     WhereNode.as_mql = where_node
-    WhereNode.as_mql_idx = where_node_idx