Merge PR #353

polymorphic/query.py

@@ -325,12 +325,12 @@ class PolymorphicQuerySet(QuerySet):
         Finally we re-sort the resulting objects into the correct order and
         return them as a list.
         """
-        ordered_id_list = []    # list of ids of result-objects in correct order
+        resultlist = []            # polymorphic list of result-objects
-        results = {}            # polymorphic dict of result-objects, keyed with their id (no order)
 
         # dict contains one entry per unique model type occurring in result,
         # in the format idlist_per_model[modelclass]=[list-of-object-ids]
         idlist_per_model = defaultdict(list)
+        indexlist_per_model = defaultdict(list)
 
         # django's automatic ".pk" field does not always work correctly for
         # custom fields in derived objects (unclear yet who to put the blame on).
@@ -342,38 +342,32 @@ class PolymorphicQuerySet(QuerySet):
         pk_name = self.model.polymorphic_primary_key_name
 
         # - sort base_result_object ids into idlist_per_model lists, depending on their real class;
-        # - also record the correct result order in "ordered_id_list"
         # - store objects that already have the correct class into "results"
-        base_result_objects_by_id = {}
         content_type_manager = ContentType.objects.db_manager(self.db)
         self_model_class_id = content_type_manager.get_for_model(self.model, for_concrete_model=False).pk
         self_concrete_model_class_id = content_type_manager.get_for_model(self.model, for_concrete_model=True).pk
 
-        for base_object in base_result_objects:
+        for i, base_object in enumerate(base_result_objects):
-            ordered_id_list.append(base_object.pk)
 
-            # check if id of the result object occurres more than once - this can happen e.g. with base_objects.extra(tables=...)
+            if base_object.polymorphic_ctype_id == self_model_class_id:
-            if base_object.pk not in base_result_objects_by_id:
+                # Real class is exactly the same as base class, go straight to results
-                base_result_objects_by_id[base_object.pk] = base_object
+                resultlist.append(base_object)
-
+            else:
-                if base_object.polymorphic_ctype_id == self_model_class_id:
+                real_concrete_class = base_object.get_real_instance_class()
-                    # Real class is exactly the same as base class, go straight to results
+                real_concrete_class_id = base_object.get_real_concrete_instance_class_id()
-                    results[base_object.pk] = base_object
 
+                if real_concrete_class_id is None:
+                    # Dealing with a stale content type
+                    continue
+                elif real_concrete_class_id == self_concrete_model_class_id:
+                    # Real and base classes share the same concrete ancestor,
+                    # upcast it and put it in the results
+                    resultlist.append(transmogrify(real_concrete_class, base_object))
                 else:
-                    real_concrete_class = base_object.get_real_instance_class()
+                    real_concrete_class = content_type_manager.get_for_id(real_concrete_class_id).model_class()
-                    real_concrete_class_id = base_object.get_real_concrete_instance_class_id()
+                    idlist_per_model[real_concrete_class].append(getattr(base_object, pk_name))
-
+                    indexlist_per_model[real_concrete_class].append((i, len(resultlist)))
-                    if real_concrete_class_id is None:
+                    resultlist.append(None)
-                        # Dealing with a stale content type
-                        continue
-                    elif real_concrete_class_id == self_concrete_model_class_id:
-                        # Real and base classes share the same concrete ancestor,
-                        # upcast it and put it in the results
-                        results[base_object.pk] = transmogrify(real_concrete_class, base_object)
-                    else:
-                        real_concrete_class = content_type_manager.get_for_id(real_concrete_class_id).model_class()
-                        idlist_per_model[real_concrete_class].append(getattr(base_object, pk_name))
 
         # For each model in "idlist_per_model" request its objects (the real model)
         # from the db and store them in results[].
@@ -381,6 +375,7 @@ class PolymorphicQuerySet(QuerySet):
         # Then we copy the extra() select fields from the base objects to the real objects.
         # TODO: defer(), only(): support for these would be around here
         for real_concrete_class, idlist in idlist_per_model.items():
+            indices = indexlist_per_model[real_concrete_class]
             real_objects = real_concrete_class._base_objects.db_manager(self.db).filter(**{
                 ('%s__in' % pk_name): idlist,
             })
@@ -413,8 +408,15 @@ class PolymorphicQuerySet(QuerySet):
                 deferred_loading_fields.append(translated_field_name)
             real_objects.query.deferred_loading = (set(deferred_loading_fields), self.query.deferred_loading[1])
 
-            for real_object in real_objects:
+            real_objects_dict = {
-                o_pk = getattr(real_object, pk_name)
+                getattr(real_object, pk_name): real_object
+                for real_object in real_objects
+            }
+
+            for i, j in indices:
+                base_object = base_result_objects[i]
+                o_pk = getattr(base_object, pk_name)
+                real_object = copy.copy(real_objects_dict[o_pk])
                 real_class = real_object.get_real_instance_class()
 
                 # If the real class is a proxy, upcast it
@@ -423,18 +425,17 @@ class PolymorphicQuerySet(QuerySet):
 
                 if self.query.annotations:
                     for anno_field_name in six.iterkeys(self.query.annotations):
-                        attr = getattr(base_result_objects_by_id[o_pk], anno_field_name)
+                        attr = getattr(base_object, anno_field_name)
                         setattr(real_object, anno_field_name, attr)
 
                 if self.query.extra_select:
                     for select_field_name in six.iterkeys(self.query.extra_select):
-                        attr = getattr(base_result_objects_by_id[o_pk], select_field_name)
+                        attr = getattr(base_object, select_field_name)
                         setattr(real_object, select_field_name, attr)
 
-                results[o_pk] = real_object
+                resultlist[j] = real_object
 
-        # re-create correct order and return result list
+        resultlist = [i for i in resultlist if i]
-        resultlist = [results[ordered_id] for ordered_id in ordered_id_list if ordered_id in results]
 
         # set polymorphic_annotate_names in all objects (currently just used for debugging/printing)
         if self.query.annotations:

polymorphic/tests/test_orm.py

@@ -73,6 +73,7 @@ from polymorphic.tests.models import (
     RelationB,
     RelationBC,
     RelationBase,
+    RelatingModel,
     RubberDuck,
     SubclassSelectorAbstractBaseModel,
     SubclassSelectorAbstractConcreteModel,
@@ -1055,6 +1056,7 @@ class PolymorphicTests(TransactionTestCase):
                 MultiTableDerived(field1='field1', field2='field2')
             ])
 
+
     def test_can_query_using_subclass_selector_on_abstract_model(self):
         obj = SubclassSelectorAbstractConcreteModel.objects.create(concrete_field='abc')
 
@@ -1072,3 +1074,12 @@ class PolymorphicTests(TransactionTestCase):
         ).get()
 
         self.assertEqual(obj.pk, queried_obj.pk)
+
+    def test_prefetch_related_behaves_normally_with_polymorphic_model(self):
+        b1 = RelatingModel.objects.create()
+        b2 = RelatingModel.objects.create()
+        a = b1.many2many.create()
+        b2.many2many.add(a)
+        qs = RelatingModel.objects.prefetch_related('many2many')
+        for obj in qs:
+            self.assertEqual(len(obj.many2many.all()), 1)