Merge 78e5ea4 into f57e37d

charettes · Jul 12, 2019 · 43d4cd3 · 43d4cd3
2 parents f57e37d + 78e5ea4
commit 43d4cd3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/polymodels/models.py b/polymodels/models.py
@@ -36,7 +36,10 @@ def __call__(self, obj, with_prefetched_objects=False):
         if proxy:
             casted = copy_fields(casted, proxy)
         if with_prefetched_objects:
-            casted._prefetched_objects_cache.update(obj._prefetched_objects_cache)
+            try:
+                casted._prefetched_objects_cache.update(obj._prefetched_objects_cache)
+            except AttributeError:
+                casted._prefetched_objects_cache = obj._prefetched_objects_cache
         return casted
 
 

diff --git a/tests/test_managers.py b/tests/test_managers.py
@@ -108,6 +108,18 @@ def test_select_subclasses_prefetch_related(self):
         zoo.animals.add(animal, mammal, monkey)
         other_monkey = Monkey.objects.create(name='monkey')
         monkey.friends.add(other_monkey)
+        queryset = Animal.objects.select_subclasses().prefetch_related('zoos')
+        with self.assertNumQueries(2):
+            self.assertSequenceEqual(queryset, [
+                animal,
+                mammal,
+                monkey,
+                other_monkey,
+            ])
+            self.assertSequenceEqual(queryset[0].zoos.all(), [zoo])
+            self.assertSequenceEqual(queryset[1].zoos.all(), [zoo])
+            self.assertSequenceEqual(queryset[2].zoos.all(), [zoo])
+        # Test prefetch related combination.
         queryset = Animal.objects.select_subclasses().prefetch_related(
             'zoos',
             'mammal__monkey__friends',