add support for CombinedExpression

Author: timgraham

.github/workflows/test-python.yml

@@ -76,6 +76,7 @@ jobs:
           datetimes
           db_functions
           empty
+          expressions.tests.ExpressionOperatorTests
           expressions.tests.NegatedExpressionTests
           defer
           defer_regress

django_mongodb/expressions.py

@@ -1,10 +1,24 @@
-from django.db.models.expressions import Col, ExpressionWrapper, NegatedExpression, Value
+from django.db.models.expressions import (
+    Col,
+    CombinedExpression,
+    ExpressionWrapper,
+    NegatedExpression,
+    Value,
+)
 
 
 def col(self, compiler, connection):  # noqa: ARG001
     return f"${self.target.column}"
 
 
+def combined_expression(self, compiler, connection):
+    expressions = [
+        self.lhs.as_mql(compiler, connection),
+        self.rhs.as_mql(compiler, connection),
+    ]
+    return connection.ops.combine_expression(self.connector, expressions)
+
+
 def expression_wrapper(self, compiler, connection):
     return self.expression.as_mql(compiler, connection)
 
@@ -19,6 +33,7 @@ def value(self, compiler, connection):  # noqa: ARG001
 
 def register_expressions():
     Col.as_mql = col
+    CombinedExpression.as_mql = combined_expression
     ExpressionWrapper.as_mql = expression_wrapper
     NegatedExpression.as_mql = negated_expression
     Value.as_mql = value

django_mongodb/features.py

@@ -176,18 +176,11 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_subquery_with_parameters",
             "lookup.tests.LookupQueryingTests.test_filter_subquery_lhs",
             # Invalid $project :: caused by :: Unknown expression $count,
+            "annotations.tests.NonAggregateAnnotationTestCase.test_combined_expression_annotation_with_aggregation",
+            "annotations.tests.NonAggregateAnnotationTestCase.test_combined_f_expression_annotation_with_aggregation",
             "annotations.tests.NonAggregateAnnotationTestCase.test_full_expression_annotation_with_aggregation",
             "annotations.tests.NonAggregateAnnotationTestCase.test_grouping_by_q_expression_annotation",
             "annotations.tests.NonAggregateAnnotationTestCase.test_q_expression_annotation_with_aggregation",
-            # CombinedExpression not implemented.
-            "annotations.tests.NonAggregateAnnotationTestCase.test_combined_annotation_commutative",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_combined_expression_annotation_with_aggregation",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_combined_f_expression_annotation_with_aggregation",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_decimal_annotation",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_defer_annotation",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_filter_decimal_annotation",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_mixed_type_annotation_numbers",
-            "annotations.tests.NonAggregateAnnotationTestCase.test_values_annotation",
             # Func not implemented.
             "annotations.tests.NonAggregateAnnotationTestCase.test_custom_functions",
             "annotations.tests.NonAggregateAnnotationTestCase.test_custom_functions_can_ref_other_functions",

django_mongodb/operations.py

@@ -4,12 +4,24 @@
 from bson.decimal128 import Decimal128
 from django.conf import settings
 from django.db.backends.base.operations import BaseDatabaseOperations
+from django.db.models.expressions import Combinable
 from django.utils import timezone
 from django.utils.regex_helper import _lazy_re_compile
 
 
 class DatabaseOperations(BaseDatabaseOperations):
     compiler_module = "django_mongodb.compiler"
+    combine_operators = {
+        Combinable.ADD: "add",
+        Combinable.SUB: "subtract",
+        Combinable.MUL: "multiply",
+        Combinable.DIV: "divide",
+        Combinable.POW: "pow",
+        Combinable.MOD: "mod",
+        Combinable.BITAND: "bitAnd",
+        Combinable.BITOR: "bitOr",
+        Combinable.BITXOR: "bitXor",
+    }
 
     def adapt_datefield_value(self, value):
         """Store DateField as datetime."""
@@ -86,6 +98,15 @@ def convert_uuidfield_value(self, value, expression, connection):
             value = uuid.UUID(value)
         return value
 
+    def combine_expression(self, connector, sub_expressions):
+        lhs, rhs = sub_expressions
+        if connector == Combinable.BITLEFTSHIFT:
+            return {"$multiply": [lhs, {"$pow": [2, rhs]}]}
+        if connector == Combinable.BITRIGHTSHIFT:
+            return {"$floor": {"$divide": [lhs, {"$pow": [2, rhs]}]}}
+        operator = self.combine_operators[connector]
+        return {f"${operator}": sub_expressions}
+
     def prep_for_like_query(self, x):
         # Override value escaping for LIKE queries.
         return str(x)