Clean up RestrictedGenericForeignKey and RestrictedPrefetch

This commit is contained in:
jeremystretch 2022-10-21 12:14:46 -04:00
parent ac33eb71ee
commit b157f6469b
2 changed files with 14 additions and 7 deletions

View File

@ -78,10 +78,18 @@ class NaturalOrderingField(models.CharField):
class RestrictedGenericForeignKey(GenericForeignKey): class RestrictedGenericForeignKey(GenericForeignKey):
# Replicated from GenericForeignKey # Replicated largely from GenericForeignKey. Changes include:
# 1. Capture restrict_params from RestrictedPrefetch (hack)
# 2. If restrict_params is set, call restrict() on the queryset for
# the related model
def get_prefetch_queryset(self, instances, queryset=None): def get_prefetch_queryset(self, instances, queryset=None):
restrict_params = {}
# Compensate for the hack in RestrictedPrefetch # Compensate for the hack in RestrictedPrefetch
restrict_params = queryset if type(queryset) is dict else {} if type(queryset) is dict:
restrict_params = queryset
elif queryset is not None:
raise ValueError("Custom queryset can't be used for this lookup.")
# For efficiency, group the instances by content type and then do one # For efficiency, group the instances by content type and then do one
# query per model # query per model

View File

@ -21,13 +21,12 @@ class RestrictedPrefetch(Prefetch):
'action': self.restrict_action, 'action': self.restrict_action,
} }
qs = super().get_current_queryset(level) if qs := super().get_current_queryset(level):
if qs: return qs.restrict(**params)
return qs.filter(**params)
# Bit of a hack. If no queryset is defined, pass through the dict of restrict() # Bit of a hack. If no queryset is defined, pass through the dict of restrict()
# kwargs to be called later. This is necessary e.g. for GenericForeignKey fields, # kwargs to be handled by the field. This is necessary e.g. for GenericForeignKey
# which do not permit setting a queryset on a Prefetch object. # fields, which do not permit setting a queryset on a Prefetch object.
return params return params