Merge PR #353
commit
3d9587acfb
|
|
@ -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
|
||||
results = {} # polymorphic dict of result-objects, keyed with their id (no order)
|
||||
resultlist = [] # polymorphic list of result-objects
|
||||
|
||||
# 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,24 +342,16 @@ 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:
|
||||
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.pk not in base_result_objects_by_id:
|
||||
base_result_objects_by_id[base_object.pk] = base_object
|
||||
for i, base_object in enumerate(base_result_objects):
|
||||
|
||||
if base_object.polymorphic_ctype_id == self_model_class_id:
|
||||
# Real class is exactly the same as base class, go straight to results
|
||||
results[base_object.pk] = base_object
|
||||
|
||||
resultlist.append(base_object)
|
||||
else:
|
||||
real_concrete_class = base_object.get_real_instance_class()
|
||||
real_concrete_class_id = base_object.get_real_concrete_instance_class_id()
|
||||
|
|
@ -370,10 +362,12 @@ class PolymorphicQuerySet(QuerySet):
|
|||
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)
|
||||
resultlist.append(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))
|
||||
indexlist_per_model[real_concrete_class].append((i, len(resultlist)))
|
||||
resultlist.append(None)
|
||||
|
||||
# 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:
|
||||
o_pk = getattr(real_object, pk_name)
|
||||
real_objects_dict = {
|
||||
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 = [results[ordered_id] for ordered_id in ordered_id_list if ordered_id in results]
|
||||
resultlist = [i for i in resultlist if i]
|
||||
|
||||
# set polymorphic_annotate_names in all objects (currently just used for debugging/printing)
|
||||
if self.query.annotations:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue