diff --git a/polymodels/models.py b/polymodels/models.py index 9689157..2e26c43 100644 --- 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 index 1798da3..0a7756d 100644 --- 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',