move aggregates to separate file

Author: timgraham

django_mongodb/__init__.py

@@ -6,12 +6,14 @@
 
 check_django_compatability()
 
+from .aggregates import register_aggregates  # noqa: E402
 from .expressions import register_expressions  # noqa: E402
 from .fields import register_fields  # noqa: E402
 from .functions import register_functions  # noqa: E402
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
 
+register_aggregates()
 register_expressions()
 register_fields()
 register_functions()

django_mongodb/aggregates.py

@@ -0,0 +1,74 @@
+from copy import deepcopy
+
+from django.db.models.aggregates import Aggregate, Count, StdDev, Variance
+from django.db.models.expressions import Case, Value, When
+from django.db.models.lookups import Exact
+from django.db.models.sql.where import WhereNode
+
+from .query_utils import process_lhs
+
+MONGO_AGGREGATIONS = {
+    Count: "sum",
+    StdDev: "stdDev",  # Samp or Pop suffix added in aggregate().
+    Variance: "stdDev",  # Likewise.
+}
+
+
+def aggregate(self, compiler, connection, **extra_context):  # noqa: ARG001
+    if self.filter:
+        node = self.copy()
+        node.filter = None
+        source_expressions = node.get_source_expressions()
+        condition = When(self.filter, then=source_expressions[0])
+        node.set_source_expressions([Case(condition)] + source_expressions[1:])
+    else:
+        node = self
+    lhs_mql = process_lhs(node, compiler, connection)
+    operator = MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
+    # Add suffixes to StdDev/Variance.
+    if self.function.endswith("_SAMP"):
+        operator += "Samp"
+    elif self.function.endswith("_POP"):
+        operator += "Pop"
+    return {f"${operator}": lhs_mql}
+
+
+def count(self, compiler, connection, resolve_inner_expression=False, **extra_context):  # noqa: ARG001
+    """
+    When resolve_inner_expression is True, return the argument as MQL that
+    resolves as a value. This is used to count different elements, so the inner
+    values are returned to be pushed into a set.
+    """
+    if not self.distinct or resolve_inner_expression:
+        if self.filter:
+            node = self.copy()
+            node.filter = None
+            source_expressions = node.get_source_expressions()
+            filter_ = deepcopy(self.filter)
+            filter_.add(
+                WhereNode([Exact(source_expressions[0], Value(None))], negated=True),
+                filter_.default,
+            )
+            condition = When(filter_, then=Value(1))
+            node.set_source_expressions([Case(condition)] + source_expressions[1:])
+            inner_expression = process_lhs(node, compiler, connection)
+        else:
+            lhs_mql = process_lhs(self, compiler, connection)
+            null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
+            inner_expression = {
+                "$cond": {"if": null_cond, "then": None, "else": lhs_mql if self.distinct else 1}
+            }
+        if resolve_inner_expression:
+            return inner_expression
+        return {"$sum": inner_expression}
+    # If distinct=True or resolve_inner_expression=False, sum the size
+    # of the set.
+    lhs_mql = process_lhs(self, compiler, connection)
+    # Subtract 1 if None is in the set (it shouldn't have been counted).
+    exits_null = {"$cond": {"if": {"$in": [{"$literal": None}, lhs_mql]}, "then": -1, "else": 0}}
+    return {"$add": [{"$size": lhs_mql}, exits_null]}
+
+
+def register_aggregates():
+    Aggregate.as_mql = aggregate
+    Count.as_mql = count

django_mongodb/expressions.py

@@ -13,6 +13,7 @@
     NegatedExpression,
     Ref,
     ResolvedOuterRef,
+    Star,
     Subquery,
     Value,
     When,
@@ -79,6 +80,10 @@ def ref(self, compiler, connection):  # noqa: ARG001
     return f"${self.refs}"
 
 
+def star(self, compiler, connection):  # noqa: ARG001
+    return {"$literal": True}
+
+
 def subquery(self, compiler, connection):  # noqa: ARG001
     raise NotSupportedError(f"{self.__class__.__name__} is not supported on MongoDB.")
 
@@ -113,6 +118,7 @@ def register_expressions():
     Query.as_mql = query
     Ref.as_mql = ref
     ResolvedOuterRef.as_mql = ResolvedOuterRef.as_sql
+    Star.as_mql = star
     Subquery.as_mql = subquery
     When.as_mql = when
     Value.as_mql = value

django_mongodb/functions.py

@@ -1,9 +1,5 @@
-from copy import deepcopy
-
 from django.db import NotSupportedError
-from django.db.models.aggregates import Aggregate, Count, StdDev, Variance
-from django.db.models.expressions import Case, Func, Star, Value, When
-from django.db.models.functions import Now
+from django.db.models.expressions import Func
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
 from django.db.models.functions.datetime import (
     Extract,
@@ -17,6 +13,7 @@
     ExtractWeek,
     ExtractWeekDay,
     ExtractYear,
+    Now,
     TruncBase,
 )
 from django.db.models.functions.math import Ceil, Cot, Degrees, Log, Power, Radians, Random, Round
@@ -34,16 +31,9 @@
     Trim,
     Upper,
 )
-from django.db.models.lookups import Exact
-from django.db.models.sql.where import WhereNode
 
 from .query_utils import process_lhs
 
-MONGO_AGGREGATIONS = {
-    Count: "sum",
-    StdDev: "stdDev",  # Samp or Pop suffix added in aggregate().
-    Variance: "stdDev",  # Likewise.
-}
 MONGO_OPERATORS = {
     Ceil: "ceil",
     Coalesce: "ifNull",
@@ -68,25 +58,6 @@
 }
 
 
-def aggregate(self, compiler, connection, **extra_context):  # noqa: ARG001
-    if self.filter:
-        node = self.copy()
-        node.filter = None
-        source_expressions = node.get_source_expressions()
-        condition = When(self.filter, then=source_expressions[0])
-        node.set_source_expressions([Case(condition)] + source_expressions[1:])
-    else:
-        node = self
-    lhs_mql = process_lhs(node, compiler, connection)
-    operator = MONGO_AGGREGATIONS.get(self.__class__, self.function.lower())
-    # Add suffixes to StdDev/Variance.
-    if self.function.endswith("_SAMP"):
-        operator += "Samp"
-    elif self.function.endswith("_POP"):
-        operator += "Pop"
-    return {f"${operator}": lhs_mql}
-
-
 def cast(self, compiler, connection):
     output_type = connection.data_types[self.output_field.get_internal_type()]
     lhs_mql = process_lhs(self, compiler, connection)[0]
@@ -117,42 +88,6 @@ def cot(self, compiler, connection):
     return {"$divide": [1, {"$tan": lhs_mql}]}
 
 
-def count(self, compiler, connection, resolve_inner_expression=False, **extra_context):  # noqa: ARG001
-    """
-    When resolve_inner_expression is True, return the argument as MQL that
-    resolves as a value. This is used to count different elements, so the inner
-    values are returned to be pushed into a set.
-    """
-    if not self.distinct or resolve_inner_expression:
-        if self.filter:
-            node = self.copy()
-            node.filter = None
-            source_expressions = node.get_source_expressions()
-            filter_ = deepcopy(self.filter)
-            filter_.add(
-                WhereNode([Exact(source_expressions[0], Value(None))], negated=True),
-                filter_.default,
-            )
-            condition = When(filter_, then=Value(1))
-            node.set_source_expressions([Case(condition)] + source_expressions[1:])
-            inner_expression = process_lhs(node, compiler, connection)
-        else:
-            lhs_mql = process_lhs(self, compiler, connection)
-            null_cond = {"$in": [{"$type": lhs_mql}, ["missing", "null"]]}
-            inner_expression = {
-                "$cond": {"if": null_cond, "then": None, "else": lhs_mql if self.distinct else 1}
-            }
-        if resolve_inner_expression:
-            return inner_expression
-        return {"$sum": inner_expression}
-    # If distinct=True or resolve_inner_expression=False, sum the size
-    # of the set.
-    lhs_mql = process_lhs(self, compiler, connection)
-    # Subtract 1 if None is in the set (it shouldn't have been counted).
-    exits_null = {"$cond": {"if": {"$in": [{"$literal": None}, lhs_mql]}, "then": -1, "else": 0}}
-    return {"$add": [{"$size": lhs_mql}, exits_null]}
-
-
 def extract(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)
     operator = EXTRACT_OPERATORS.get(self.lookup_name)
@@ -223,10 +158,6 @@ def round_(self, compiler, connection):
     return {"$round": [expr.as_mql(compiler, connection) for expr in self.get_source_expressions()]}
 
 
-def star(self, compiler, connection):  # noqa: ARG001
-    return {"$literal": True}
-
-
 def str_index(self, compiler, connection):
     lhs = process_lhs(self, compiler, connection)
     # StrIndex should be 0-indexed (not found) but it's -1-indexed on MongoDB.
@@ -261,12 +192,10 @@ def trunc(self, compiler, connection):
 
 
 def register_functions():
-    Aggregate.as_mql = aggregate
     Cast.as_mql = cast
     Concat.as_mql = concat
     ConcatPair.as_mql = concat_pair
     Cot.as_mql = cot
-    Count.as_mql = count
     Extract.as_mql = extract
     Func.as_mql = func
     Left.as_mql = left
@@ -279,7 +208,6 @@ def register_functions():
     Replace.as_mql = replace
     Round.as_mql = round_
     RTrim.as_mql = trim("rtrim")
-    Star.as_mql = star
     StrIndex.as_mql = str_index
     Substr.as_mql = substr
     Trim.as_mql = trim("trim")