Fixed #16436 -- defer + annotate + select_related crash

Correctly calculate the ``aggregate_start`` offset from loaded fields,
if any are deferred, instead of ``self.query.select`` which includes all
fields on the model.

Also made some PEP 8 fixes.

1 files changed, 34 insertions, 7 deletions

diff --git a/tests/defer_regress/tests.py b/tests/defer_regress/tests.py
index fc77793288..ad2546794c 100644
--- a/tests/defer_regress/tests.py
+++ b/tests/defer_regress/tests.py
@@ -8,8 +8,9 @@ from django.db.models import Count
 from django.db.models.loading import cache
 from django.test import TestCase
 
-from .models import (ResolveThis, Item, RelatedItem, Child, Leaf, Proxy,
-    SimpleItem, Feature, ItemAndSimpleItem, OneToOneItem, SpecialFeature)
+from .models import (
+    ResolveThis, Item, RelatedItem, Child, Leaf, Proxy, SimpleItem, Feature,
+    ItemAndSimpleItem, OneToOneItem, SpecialFeature, Location, Request)
 
 
 class DeferRegressionTest(TestCase):
@@ -76,7 +77,9 @@ class DeferRegressionTest(TestCase):
         self.assertEqual(results[0].child.name, "c1")
         self.assertEqual(results[0].second_child.name, "c2")
 
-        results = Leaf.objects.only("name", "child", "second_child", "child__name", "second_child__name").select_related()
+        results = Leaf.objects.only(
+            "name", "child", "second_child", "child__name", "second_child__name"
+        ).select_related()
         self.assertEqual(results[0].child.name, "c1")
         self.assertEqual(results[0].second_child.name, "c2")
 
@@ -144,6 +147,14 @@ class DeferRegressionTest(TestCase):
                 cache.get_app("defer_regress"), include_deferred=True))
         )
 
+        # Regression for #16409 - make sure defer() and only() work with annotate()
+        self.assertIsInstance(
+            list(SimpleItem.objects.annotate(Count('feature')).defer('name')),
+            list)
+        self.assertIsInstance(
+            list(SimpleItem.objects.annotate(Count('feature')).only('name')),
+            list)
+
     def test_only_and_defer_usage_on_proxy_models(self):
         # Regression for #15790 - only() broken for proxy models
         proxy = Proxy.objects.create(name="proxy", value=42)
@@ -160,7 +171,7 @@ class DeferRegressionTest(TestCase):
         self.assertEqual(dp.value, proxy.value, msg=msg)
 
     def test_resolve_columns(self):
-        rt = ResolveThis.objects.create(num=5.0, name='Foobar')
+        ResolveThis.objects.create(num=5.0, name='Foobar')
         qs = ResolveThis.objects.defer('num')
         self.assertEqual(1, qs.count())
         self.assertEqual('Foobar', qs[0].name)
@@ -196,7 +207,7 @@ class DeferRegressionTest(TestCase):
         item1 = Item.objects.create(name="first", value=47)
         item2 = Item.objects.create(name="second", value=42)
         simple = SimpleItem.objects.create(name="simple", value="23")
-        related = ItemAndSimpleItem.objects.create(item=item1, simple=simple)
+        ItemAndSimpleItem.objects.create(item=item1, simple=simple)
 
         obj = ItemAndSimpleItem.objects.defer('item').select_related('simple').get()
         self.assertEqual(obj.item, item1)
@@ -223,7 +234,23 @@ class DeferRegressionTest(TestCase):
 
     def test_deferred_class_factory(self):
         from django.db.models.query_utils import deferred_class_factory
-        new_class = deferred_class_factory(Item,
+        new_class = deferred_class_factory(
+            Item,
             ('this_is_some_very_long_attribute_name_so_modelname_truncation_is_triggered',))
-        self.assertEqual(new_class.__name__,
+        self.assertEqual(
+            new_class.__name__,
             'Item_Deferred_this_is_some_very_long_attribute_nac34b1f495507dad6b02e2cb235c875e')
+
+class DeferAnnotateSelectRelatedTest(TestCase):
+    def test_defer_annotate_select_related(self):
+        location = Location.objects.create()
+        Request.objects.create(location=location)
+        self.assertIsInstance(list(Request.objects
+            .annotate(Count('items')).select_related('profile', 'location')
+            .only('profile', 'location')), list)
+        self.assertIsInstance(list(Request.objects
+            .annotate(Count('items')).select_related('profile', 'location')
+            .only('profile__profile1', 'location__location1')), list)
+        self.assertIsInstance(list(Request.objects
+            .annotate(Count('items')).select_related('profile', 'location')
+            .defer('request1', 'request2', 'request3', 'request4')), list)