AfterFetchQuerySetMixin for Django

after_fetch_queryset_mixin.py

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.