Fixed django#373 - Add Count support.

django/db/models/aggregates.py

@@ -3,7 +3,8 @@
 """
 
 from django.core.exceptions import FieldError, FullResultSet
-from django.db.models.expressions import Case, Func, Star, Value, When
+from django.db import NotSupportedError
+from django.db.models.expressions import Case, ColPairs, Func, Star, Value, When
 from django.db.models.fields import IntegerField
 from django.db.models.functions.comparison import Coalesce
 from django.db.models.functions.mixins import (
@@ -174,6 +175,29 @@ def __init__(self, expression, filter=None, **extra):
             raise ValueError("Star cannot be used with filter. Please specify a field.")
         super().__init__(expression, filter=filter, **extra)
 
+    def resolve_expression(self, *args, **kwargs):
+        result = super().resolve_expression(*args, **kwargs)
+        expr = result.source_expressions[0]
+
+        # In case of composite primary keys, count the first column.
+        if isinstance(expr, ColPairs):
+            if self.distinct:
+                raise NotSupportedError(
+                    "COUNT(DISTINCT) doesn't support composite primary keys"
+                )
+
+            cols = expr.get_cols()
+            fields = [col.field for col in cols]
+            pk_fields = expr.field.model._meta.pk_fields
+            assert fields == pk_fields, (
+                "Count doesn't support counting composite fields that are not composite"
+                "primary keys"
+            )
+
+            return Count(cols[0], filter=result.filter)
+
+        return result
+
 
 class Max(Aggregate):
     function = "MAX"

django/db/models/fields/composite.py

@@ -139,8 +139,12 @@ def _check_field_name(self):
 
 
 def unnest(fields):
+    result = []
+
     for field in fields:
         if isinstance(field, CompositePrimaryKey):
-            yield from field.fields
+            result.extend(field.fields)
         else:
-            yield field
+            result.append(field)
+
+    return result

docs/howto/composite-primary-key.txt

@@ -132,18 +132,16 @@ Many database functions only accept a single field.
 
 .. code-block:: sql
 
-    COUNT("order_id")  -- OK
-    COUNT("product_id", "order_id")  -- ERROR
+    MAX("order_id")  -- OK
+    MAX("product_id", "order_id")  -- ERROR
 
 These database functions cannot be used with composite primary keys, as
 composite primary keys are a proxy to multiple fields.
 
 .. code-block:: pycon
 
-    Count("order_id")  -- OK
-    Count("pk")  -- ERROR
-    Count("foo_set__id")  -- OK
-    Count("foo_set")  -- ERROR
+    Max("order_id")  -- OK
+    Max("pk")  -- ERROR
 
 Composite primary keys and admin
 ================================

tests/composite_pk/test_aggregate.py

@@ -1,3 +1,4 @@
+from django.db import NotSupportedError
 from django.db.models import Count, Q
 from django.test import TestCase
 
@@ -55,6 +56,38 @@ def test_users_annotated_with_aliased_comments_id_count(self):
         self.assertEqual(user_3, self.user_3)
         self.assertEqual(user_3.comments_count, 3)
 
+    def test_users_annotated_with_comments_count(self):
+        user_1, user_2, user_3 = User.objects.annotate(Count("comments")).order_by("pk")
+
+        self.assertEqual(user_1, self.user_1)
+        self.assertEqual(user_1.comments__count, 2)
+        self.assertEqual(user_2, self.user_2)
+        self.assertEqual(user_2.comments__count, 1)
+        self.assertEqual(user_3, self.user_3)
+        self.assertEqual(user_3.comments__count, 3)
+
+    def test_users_annotated_with_comments_count_filter(self):
+        user_1, user_2, user_3 = User.objects.annotate(
+            comments__count=Count(
+                "comments", filter=Q(pk__in=[self.user_1.pk, self.user_2.pk])
+            )
+        ).order_by("pk")
+
+        self.assertEqual(user_1, self.user_1)
+        self.assertEqual(user_1.comments__count, 2)
+        self.assertEqual(user_2, self.user_2)
+        self.assertEqual(user_2.comments__count, 1)
+        self.assertEqual(user_3, self.user_3)
+        self.assertEqual(user_3.comments__count, 0)
+
+    def test_count_distinct_not_supported(self):
+        with self.assertRaisesMessage(
+            NotSupportedError, "COUNT(DISTINCT) doesn't support composite primary keys"
+        ):
+            self.assertIsNone(
+                User.objects.annotate(comments__count=Count("comments", distinct=True))
+            )
+
     def test_user_values_annotated_with_comments_id_count(self):
         self.assertSequenceEqual(
             User.objects.values("pk").annotate(Count("comments__id")).order_by("pk"),