Fixed #27849 -- Added filtering support to aggregates.

3 files changed, 104 insertions, 1 deletions

diff --git a/tests/aggregation/test_filter_argument.py b/tests/aggregation/test_filter_argument.py
new file mode 100644
index 0000000000..54836178c4
--- /dev/null
+++ b/tests/aggregation/test_filter_argument.py
@@ -0,0 +1,81 @@
+import datetime
+from decimal import Decimal
+
+from django.db.models import Case, Count, F, Q, Sum, When
+from django.test import TestCase
+
+from .models import Author, Book, Publisher
+
+
+class FilteredAggregateTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.a1 = Author.objects.create(name='test', age=40)
+        cls.a2 = Author.objects.create(name='test2', age=60)
+        cls.a3 = Author.objects.create(name='test3', age=100)
+        cls.p1 = Publisher.objects.create(name='Apress', num_awards=3, duration=datetime.timedelta(days=1))
+        cls.b1 = Book.objects.create(
+            isbn='159059725', name='The Definitive Guide to Django: Web Development Done Right',
+            pages=447, rating=4.5, price=Decimal('30.00'), contact=cls.a1, publisher=cls.p1,
+            pubdate=datetime.date(2007, 12, 6),
+        )
+        cls.b2 = Book.objects.create(
+            isbn='067232959', name='Sams Teach Yourself Django in 24 Hours',
+            pages=528, rating=3.0, price=Decimal('23.09'), contact=cls.a2, publisher=cls.p1,
+            pubdate=datetime.date(2008, 3, 3),
+        )
+        cls.b3 = Book.objects.create(
+            isbn='159059996', name='Practical Django Projects',
+            pages=600, rating=4.5, price=Decimal('29.69'), contact=cls.a3, publisher=cls.p1,
+            pubdate=datetime.date(2008, 6, 23),
+        )
+        cls.a1.friends.add(cls.a2)
+        cls.a1.friends.add(cls.a3)
+        cls.b1.authors.add(cls.a1)
+        cls.b1.authors.add(cls.a3)
+        cls.b2.authors.add(cls.a2)
+        cls.b3.authors.add(cls.a3)
+
+    def test_filtered_aggregates(self):
+        agg = Sum('age', filter=Q(name__startswith='test'))
+        self.assertEqual(Author.objects.aggregate(age=agg)['age'], 200)
+
+    def test_double_filtered_aggregates(self):
+        agg = Sum('age', filter=Q(Q(name='test2') & ~Q(name='test')))
+        self.assertEqual(Author.objects.aggregate(age=agg)['age'], 60)
+
+    def test_excluded_aggregates(self):
+        agg = Sum('age', filter=~Q(name='test2'))
+        self.assertEqual(Author.objects.aggregate(age=agg)['age'], 140)
+
+    def test_related_aggregates_m2m(self):
+        agg = Sum('friends__age', filter=~Q(friends__name='test'))
+        self.assertEqual(Author.objects.filter(name='test').aggregate(age=agg)['age'], 160)
+
+    def test_related_aggregates_m2m_and_fk(self):
+        q = Q(friends__book__publisher__name='Apress') & ~Q(friends__name='test3')
+        agg = Sum('friends__book__pages', filter=q)
+        self.assertEqual(Author.objects.filter(name='test').aggregate(pages=agg)['pages'], 528)
+
+    def test_plain_annotate(self):
+        agg = Sum('book__pages', filter=Q(book__rating__gt=3))
+        qs = Author.objects.annotate(pages=agg).order_by('pk')
+        self.assertSequenceEqual([a.pages for a in qs], [447, None, 1047])
+
+    def test_filtered_aggregate_on_annotate(self):
+        pages_annotate = Sum('book__pages', filter=Q(book__rating__gt=3))
+        age_agg = Sum('age', filter=Q(total_pages__gte=400))
+        aggregated = Author.objects.annotate(total_pages=pages_annotate).aggregate(summed_age=age_agg)
+        self.assertEqual(aggregated, {'summed_age': 140})
+
+    def test_case_aggregate(self):
+        agg = Sum(
+            Case(When(friends__age=40, then=F('friends__age'))),
+            filter=Q(friends__name__startswith='test'),
+        )
+        self.assertEqual(Author.objects.aggregate(age=agg)['age'], 80)
+
+    def test_sum_star_exception(self):
+        msg = 'Star cannot be used with filter. Please specify a field.'
+        with self.assertRaisesMessage(ValueError, msg):
+            Count('*', filter=Q(age=40))
diff --git a/tests/expressions/tests.py b/tests/expressions/tests.py
index 8952045002..7865492ed5 100644
--- a/tests/expressions/tests.py
+++ b/tests/expressions/tests.py
@@ -5,7 +5,7 @@ from copy import deepcopy
 
 from django.core.exceptions import FieldError
 from django.db import DatabaseError, connection, models, transaction
-from django.db.models import CharField, TimeField, UUIDField
+from django.db.models import CharField, Q, TimeField, UUIDField
 from django.db.models.aggregates import (
     Avg, Count, Max, Min, StdDev, Sum, Variance,
 )
@@ -1369,3 +1369,16 @@ class ReprTests(TestCase):
         self.assertEqual(repr(StdDev('a')), "StdDev(F(a), sample=False)")
         self.assertEqual(repr(Sum('a')), "Sum(F(a))")
         self.assertEqual(repr(Variance('a', sample=True)), "Variance(F(a), sample=True)")
+
+    def test_filtered_aggregates(self):
+        filter = Q(a=1)
+        self.assertEqual(repr(Avg('a', filter=filter)), "Avg(F(a), filter=(AND: ('a', 1)))")
+        self.assertEqual(repr(Count('a', filter=filter)), "Count(F(a), distinct=False, filter=(AND: ('a', 1)))")
+        self.assertEqual(repr(Max('a', filter=filter)), "Max(F(a), filter=(AND: ('a', 1)))")
+        self.assertEqual(repr(Min('a', filter=filter)), "Min(F(a), filter=(AND: ('a', 1)))")
+        self.assertEqual(repr(StdDev('a', filter=filter)), "StdDev(F(a), filter=(AND: ('a', 1)), sample=False)")
+        self.assertEqual(repr(Sum('a', filter=filter)), "Sum(F(a), filter=(AND: ('a', 1)))")
+        self.assertEqual(
+            repr(Variance('a', sample=True, filter=filter)),
+            "Variance(F(a), filter=(AND: ('a', 1)), sample=True)"
+        )
diff --git a/tests/expressions_case/tests.py b/tests/expressions_case/tests.py
index 20d1e801ec..090607e8b8 100644
--- a/tests/expressions_case/tests.py
+++ b/tests/expressions_case/tests.py
@@ -1255,6 +1255,15 @@ class CaseDocumentationExamples(TestCase):
         )
         self.assertEqual(
             Client.objects.aggregate(
+                regular=models.Count('pk', filter=Q(account_type=Client.REGULAR)),
+                gold=models.Count('pk', filter=Q(account_type=Client.GOLD)),
+                platinum=models.Count('pk', filter=Q(account_type=Client.PLATINUM)),
+            ),
+            {'regular': 2, 'gold': 1, 'platinum': 3}
+        )
+        # This was the example before the filter argument was added.
+        self.assertEqual(
+            Client.objects.aggregate(
                 regular=models.Sum(Case(
                     When(account_type=Client.REGULAR, then=1),
                     output_field=models.IntegerField(),