From 008f514d1125c821bc81e6c7fe2c14625e77c97d Mon Sep 17 00:00:00 2001
From: Raphael Michel <michel@rami.io>
Date: Tue, 6 Oct 2026 22:20:56 +0200
Subject: [PATCH] Add support for arithmetics of aggregates

Fixes #11
---
 django_ormql/expressions.py | 31 +++++++++++++++++++++++++++++++
 django_ormql/query.py       | 10 +++++++---
 tests/test_aggregate.py     | 26 ++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/django_ormql/expressions.py b/django_ormql/expressions.py
index 3205a26..ffe496f 100644
--- a/django_ormql/expressions.py
+++ b/django_ormql/expressions.py
@@ -39,6 +39,15 @@ def to_name(self, expression):
     def to_sql_type(self, expression, recurse, **kwargs):
         return self.sql_type
 
+    def is_aggregate(self, expression, recurse, **kwargs):
+        for args in expression.args.values():
+            if isinstance(args, list):
+                if any(recurse(a) for a in args):
+                    return True
+            if recurse(args):
+                return True
+        return False
+
 
 class FuncExpression(Expression):
     node_class = expressions.Anonymous
@@ -77,6 +86,13 @@ def expression_to_sql_type(expression, **kwargs):
     return None
 
 
+def expression_is_aggregate(expression, **kwargs):
+    for e in _expressions:
+        if e.matches(expression):
+            return e.is_aggregate(expression, expression_is_aggregate, **kwargs)
+    return None
+
+
 types = {
     expressions.DataType.Type.BIGDECIMAL: models.DecimalField(
         max_digits=20, decimal_places=2
@@ -292,6 +308,9 @@ def to_sql_type(self, expression, recurse, **kwargs):
             .get_internal_type()
         ]
 
+    def is_aggregate(self, expression, recurse, **kwargs):
+        return False  # barrier between query and subquery
+
 
 @register
 class Select(Expression):
@@ -326,6 +345,9 @@ def to_django(self, expression, recurse, **kwargs):
             qs,
         )
 
+    def is_aggregate(self, expression, recurse, **kwargs):
+        return False  # barrier between query and subquery
+
 
 @register
 class Alias(Expression):
@@ -795,6 +817,9 @@ def to_django(self, expression, recurse, **kwargs):
             filter=recurse(expression.expression.this, **kwargs),
         )
 
+    def is_aggregate(self, expression, recurse, **kwargs):
+        return True
+
 
 @register
 class Aggregate(Expression):
@@ -817,6 +842,9 @@ def to_django(self, expression, recurse, **kwargs):
         cls = aggregate_nodes[type(expression)]
         return cls(*args, distinct=distinct)
 
+    def is_aggregate(self, expression, recurse, **kwargs):
+        return True
+
 
 math_binary_nodes = {
     expressions.Mul: db_func.Mul,
@@ -1189,3 +1217,6 @@ def to_django(self, expression, recurse, **kwargs):
             expression_to_sql_type(expression.expressions[0], **kwargs),
             output_field=models.TextField(null=True),
         )
+
+    def is_aggregate(self, expression, recurse, **kwargs):
+        return False
diff --git a/django_ormql/query.py b/django_ormql/query.py
index 5152b3b..02358f5 100644
--- a/django_ormql/query.py
+++ b/django_ormql/query.py
@@ -7,7 +7,6 @@
     F,
     OrderBy,
     Value,
-    aggregates,
 )
 from sqlglot import (
     Dialect,
@@ -21,7 +20,12 @@
 from sqlglot.errors import ANSI_RESET, ANSI_UNDERLINE
 
 from .exceptions import QueryError, QueryNotSupported
-from .expressions import _describe_expression, expression_to_django, expression_to_name
+from .expressions import (
+    _describe_expression,
+    expression_is_aggregate,
+    expression_to_django,
+    expression_to_name,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -208,7 +212,7 @@ def _select_to_qs(self, root, parent_table_stack):
                     aggregate_names=[],
                     parent_table_stack=parent_table_stack,
                 )
-                if isinstance(django_e, aggregates.Aggregate):
+                if expression_is_aggregate(e):
                     # We do not use the alias names given by the user, first to ensure uniqueness, but also Django has
                     # had some SQL injection vulns recently that affected user-chosen annotate targets. We'll remap
                     # ourselves later.
diff --git a/tests/test_aggregate.py b/tests/test_aggregate.py
index 089adeb..1a1e186 100644
--- a/tests/test_aggregate.py
+++ b/tests/test_aggregate.py
@@ -127,3 +127,29 @@ def test_aggregate_filter(engine_t1):
         """
     )
     assert list(res) == [{"paid": 2, "canceled": 3, "all": 5}]
+
+
+@pytest.mark.django_db
+def test_aggregate_function_arithmetics(engine_t1):
+    res = engine_t1.query(
+        """
+        SELECT SUM(price) / COUNT(*) AS average
+        FROM products
+        """
+    )
+    assert [{k: round(v, 2) for k, v in row.items()} for row in res] == [
+        {
+            "average": Decimal("17.03"),
+        }
+    ]
+    res = engine_t1.query(
+        """
+        SELECT ROUND(SUM(price) / COUNT(*)) AS average
+        FROM products
+        """
+    )
+    assert [{k: round(v, 2) for k, v in row.items()} for row in res] == [
+        {
+            "average": Decimal(17),
+        }
+    ]