load_generic_related.py

class GenericRelationLoader(Manager):
    """
    Manager mixin providing a method to pre-load objects from a generic
    relationship with only in one (or two) extra queries.

    The second query is for the content type and is cached.

    Use this manager on the model with the generic foreign key and overload
    the class attributes `ct_field` and `fk_field` if necessary.

    To change the field name, where the pre-loaded objects are stored
    overload `cache_field_name`.

    The `select_related` attribute is added to the queryset for the objects of
    the generic relationship.
    """
    ct_field = 'content_type'
    fk_field = 'object_id'
    cache_field_name = '_gr_cache'
    select_related = None
    order_by = None

    def load_generic_related(self, queryset):
        """
        Add object level cache too all objects returned by `queryset`.
        Returns a list of objects.

        The pre-loaded/cached objects will be stored in a dictionary on the
        attribute with the name configured by `cache_field_name` on each of
        the objects returned.

        :param queryset: queryset for objects of a model used in the
            generic relationship of this manager's model
        :return: list of objects, result of `queryset` with object-level
            cache filled

        An example:

            users = Users.objects.all()
            users = ImageAssignment.objects.load_generic_related(users)

            # each object now has a pre-loaded generic relationship cache, eg:
            users[0]._gr_cache

        """
        # that's an internally cached query
        content_type = ContentType.objects.get_for_model(queryset.model)

        # evaluate queryset for target objects, store in pk dict
        objects = {}
        for obj in queryset.all():
            # initialize cache attribute for all
            setattr(obj, self.cache_field_name, defaultdict(list))
            objects[obj.pk] = obj

        # filter assignments by generic relationship keys
        filter = {
            self.fk_field + '__in': queryset.values_list('id', flat=True),
            self.ct_field: content_type,
        }
        gr_query = self.get_queryset().filter(**filter)
        # optionally query by another directly related model of the attachment
        if self.select_related:
            gr_query = gr_query.select_related(self.select_related)
        # optionally apply ordering
        if self.order_by:
            gr_query = gr_query.order_by(*self.order_by)

        # evaluate generic related queryset,
        # attach results to target objects in a dictionary
        for gr_object in gr_query.all():
            obj = objects[getattr(gr_object, self.fk_field)]
            cache = getattr(obj, self.cache_field_name)
            cache[gr_object.model_name].append(gr_object)

        return objects.values()