Allowed custom querysets when prefetching single valued relations

The original patch for custom prefetches didn't allow usage of custom
queryset for single valued relations (along ForeignKey or OneToOneKey).
Allowing these enables calling performance oriented queryset methods like
select_related or defer/only.

Thanks @akaariai and @timgraham for the reviews. Refs #17001.
This commit is contained in:
Loic Bistuer 2014-01-17 01:24:39 +07:00 committed by Anssi Kääriäinen
parent c8d61fa109
commit 7bbb6958dc
5 changed files with 112 additions and 23 deletions

View file

@ -1667,8 +1667,8 @@ class Prefetch(object):
def get_current_to_attr(self, level):
parts = self.prefetch_to.split(LOOKUP_SEP)
to_attr = parts[level]
to_list = self.to_attr and level == len(parts) - 1
return to_attr, to_list
as_attr = self.to_attr and level == len(parts) - 1
return to_attr, as_attr
def get_current_queryset(self, level):
if self.get_current_prefetch_to(level) == self.prefetch_to:
@ -1913,12 +1913,13 @@ def prefetch_one_level(instances, prefetcher, lookup, level):
for obj in instances:
instance_attr_val = instance_attr(obj)
vals = rel_obj_cache.get(instance_attr_val, [])
to_attr, as_attr = lookup.get_current_to_attr(level)
if single:
# Need to assign to single cache on instance
setattr(obj, cache_name, vals[0] if vals else None)
val = vals[0] if vals else None
to_attr = to_attr if as_attr else cache_name
setattr(obj, to_attr, val)
else:
to_attr, to_list = lookup.get_current_to_attr(level)
if to_list:
if as_attr:
setattr(obj, to_attr, vals)
else:
# Cache in the QuerySet.all().