kodekracker · October 17, 2022 16:02 · Sep 2, 2021
diff --git a/after_fetch_queryset_mixin.py b/after_fetch_queryset_mixin.py
@@ -0,0 +1,65 @@
+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.