class AfterFetchQuerySetMixin: """ QuerySet mixin to enable functions to run immediately after records have been fetched from the DB. """ # This is most useful for registering 'prefetch_related' like operations # or complex aggregations that need to be run after fetching, but while # still allowing chaining of other QuerySet methods. def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._after_fetch_callbacks = [] def register_after_fetch_callback(self, callback): """ Register a callback to be run after the QuerySet is fetched. The callback should be a callable that accepts a list of model instances. """ self._after_fetch_callbacks.append(callback) return self # _fetch_all and _clone are Django internals. def _fetch_all(self): already_run = self._result_cache is not None # This super() call fills out the result cache in the QuerySet, and does # any prefetches. super()._fetch_all() if already_run: # We only run our callbacks once return # Now we run our callback. for c in self._after_fetch_callbacks: c(self._result_cache) def _clone(self): retval = super()._clone() retval._after_fetch_callbacks = self._after_fetch_callbacks[:] return retval # Usage would be like this: # Example for demo purposes, there are other ways of doing this: # Suppose we want to decorate each user with an `nth_user_joined` # attribute which is calculated relative to `date_joined` attribute, # only for the batch of records retrieved. class UserQuerySet(AfterFetchQuerySet, models.QuerySet): def with_nth_user_joined(self): def add_nth_user_joined(user_list): for i, user in enumerate(sorted(user_list, key=lambda user: user.date_joined), 1): user.nth_user_joined = i return self.register_after_fetch_callback(add_nth_user_joined) # We can now do the following, with the decoration applied after the query is executed, # where that query includes the filter. users = list(User.objects.all().with_nth_user_joined().filter(id__lt=500)) # This technique is especially useful: # - if you want to do subsequent queries based on the returned values of the first query. # e.g. things like aggregations or complex prefetches. # - if you are in some framework where you don't have fully control over when the # first main query will eventually be executed, but need something to happen immediately # after that evaluation.