Fixed #25534, Fixed #31639 -- Added support for transform references in expressions.

Thanks Mariusz Felisiak and Simon Charette for reviews.

1 files changed, 99 insertions, 10 deletions

diff --git a/tests/annotations/tests.py b/tests/annotations/tests.py
index f8346d801e..cc9bd82656 100644
--- a/tests/annotations/tests.py
+++ b/tests/annotations/tests.py
@@ -4,13 +4,16 @@ from decimal import Decimal
 from django.core.exceptions import FieldDoesNotExist, FieldError
 from django.db import connection
 from django.db.models import (
-    BooleanField, Case, Count, DateTimeField, Exists, ExpressionWrapper, F,
-    FloatField, Func, IntegerField, Max, NullBooleanField, OuterRef, Q,
-    Subquery, Sum, Value, When,
+    BooleanField, Case, CharField, Count, DateTimeField, DecimalField, Exists,
+    ExpressionWrapper, F, FloatField, Func, IntegerField, Max,
+    NullBooleanField, OuterRef, Q, Subquery, Sum, Value, When,
 )
 from django.db.models.expressions import RawSQL
-from django.db.models.functions import Coalesce, ExtractYear, Length, Lower
+from django.db.models.functions import (
+    Coalesce, ExtractYear, Floor, Length, Lower, Trim,
+)
 from django.test import TestCase, skipUnlessDBFeature
+from django.test.utils import register_lookup
 
 from .models import (
     Author, Book, Company, DepartmentStore, Employee, Publisher, Store, Ticket,
@@ -95,24 +98,24 @@ class NonAggregateAnnotationTestCase(TestCase):
         cls.b5.authors.add(cls.a8, cls.a9)
         cls.b6.authors.add(cls.a8)
 
-        s1 = Store.objects.create(
+        cls.s1 = Store.objects.create(
             name='Amazon.com',
             original_opening=datetime.datetime(1994, 4, 23, 9, 17, 42),
             friday_night_closing=datetime.time(23, 59, 59)
         )
-        s2 = Store.objects.create(
+        cls.s2 = Store.objects.create(
             name='Books.com',
             original_opening=datetime.datetime(2001, 3, 15, 11, 23, 37),
             friday_night_closing=datetime.time(23, 59, 59)
         )
-        s3 = Store.objects.create(
+        cls.s3 = Store.objects.create(
             name="Mamma and Pappa's Books",
             original_opening=datetime.datetime(1945, 4, 25, 16, 24, 14),
             friday_night_closing=datetime.time(21, 30)
         )
-        s1.books.add(cls.b1, cls.b2, cls.b3, cls.b4, cls.b5, cls.b6)
-        s2.books.add(cls.b1, cls.b3, cls.b5, cls.b6)
-        s3.books.add(cls.b3, cls.b4, cls.b6)
+        cls.s1.books.add(cls.b1, cls.b2, cls.b3, cls.b4, cls.b5, cls.b6)
+        cls.s2.books.add(cls.b1, cls.b3, cls.b5, cls.b6)
+        cls.s3.books.add(cls.b3, cls.b4, cls.b6)
 
     def test_basic_annotation(self):
         books = Book.objects.annotate(is_book=Value(1))
@@ -130,6 +133,66 @@ class NonAggregateAnnotationTestCase(TestCase):
         for book in books:
             self.assertEqual(book.num_awards, book.publisher.num_awards)
 
+    def test_joined_transformed_annotation(self):
+        Employee.objects.bulk_create([
+            Employee(
+                first_name='John',
+                last_name='Doe',
+                age=18,
+                store=self.s1,
+                salary=15000,
+            ),
+            Employee(
+                first_name='Jane',
+                last_name='Jones',
+                age=30,
+                store=self.s2,
+                salary=30000,
+            ),
+            Employee(
+                first_name='Jo',
+                last_name='Smith',
+                age=55,
+                store=self.s3,
+                salary=50000,
+            ),
+        ])
+        employees = Employee.objects.annotate(
+            store_opened_year=F('store__original_opening__year'),
+        )
+        for employee in employees:
+            self.assertEqual(
+                employee.store_opened_year,
+                employee.store.original_opening.year,
+            )
+
+    def test_custom_transform_annotation(self):
+        with register_lookup(DecimalField, Floor):
+            books = Book.objects.annotate(floor_price=F('price__floor'))
+
+        self.assertSequenceEqual(books.values_list('pk', 'floor_price'), [
+            (self.b1.pk, 30),
+            (self.b2.pk, 23),
+            (self.b3.pk, 29),
+            (self.b4.pk, 29),
+            (self.b5.pk, 82),
+            (self.b6.pk, 75),
+        ])
+
+    def test_chaining_transforms(self):
+        Company.objects.create(name=' Django Software Foundation  ')
+        Company.objects.create(name='Yahoo')
+        with register_lookup(CharField, Trim), register_lookup(CharField, Length):
+            for expr in [Length('name__trim'), F('name__trim__length')]:
+                with self.subTest(expr=expr):
+                    self.assertCountEqual(
+                        Company.objects.annotate(length=expr).values('name', 'length'),
+                        [
+                            {'name': ' Django Software Foundation  ', 'length': 26},
+                            {'name': 'Yahoo', 'length': 5},
+                        ],
+                    )
+
     def test_mixed_type_annotation_date_interval(self):
         active = datetime.datetime(2015, 3, 20, 14, 0, 0)
         duration = datetime.timedelta(hours=1)
@@ -689,6 +752,23 @@ class NonAggregateAnnotationTestCase(TestCase):
             {'pub_year': 2008, 'top_rating': 4.0, 'total_pages': 1178},
         ])
 
+    def test_annotation_subquery_outerref_transform(self):
+        qs = Book.objects.annotate(
+            top_rating_year=Subquery(
+                Book.objects.filter(
+                    pubdate__year=OuterRef('pubdate__year')
+                ).order_by('-rating').values('rating')[:1]
+            ),
+        ).values('pubdate__year', 'top_rating_year')
+        self.assertCountEqual(qs, [
+            {'pubdate__year': 1991, 'top_rating_year': 5.0},
+            {'pubdate__year': 1995, 'top_rating_year': 4.0},
+            {'pubdate__year': 2007, 'top_rating_year': 4.5},
+            {'pubdate__year': 2008, 'top_rating_year': 4.0},
+            {'pubdate__year': 2008, 'top_rating_year': 4.0},
+            {'pubdate__year': 2008, 'top_rating_year': 4.0},
+        ])
+
     def test_annotation_aggregate_with_m2o(self):
         if connection.vendor == 'mysql' and 'ONLY_FULL_GROUP_BY' in connection.sql_mode:
             self.skipTest(
@@ -776,6 +856,15 @@ class AliasTests(TestCase):
             with self.subTest(book=book):
                 self.assertEqual(book.another_rating, book.rating)
 
+    def test_basic_alias_f_transform_annotation(self):
+        qs = Book.objects.alias(
+            pubdate_alias=F('pubdate'),
+        ).annotate(pubdate_year=F('pubdate_alias__year'))
+        self.assertIs(hasattr(qs.first(), 'pubdate_alias'), False)
+        for book in qs:
+            with self.subTest(book=book):
+                self.assertEqual(book.pubdate_year, book.pubdate.year)
+
     def test_alias_after_annotation(self):
         qs = Book.objects.annotate(
             is_book=Value(1),