Django relation manager for querysets

relatedmanager.py

from django.db import models


class SelectPrefetchManager(models.Manager):
    """
    This allows you to initialize a manager with 'select' or 'prefetch'
    parameters whose values are tuples of fields to be automatically applied
    to the queryset via select_related() or prefetch_related(), respectively.

    objects = RelatedManager(select=('field1',), prefetch=('m2mfield',))
    """
    def __init__(self, *args, **kwargs):
        self._select = kwargs.pop('select', None)
        self._prefetch = kwargs.pop('prefetch', None)
        super(RelatedManager, self).__init__(*args, **kwargs)

    def get_query_set(self, *args, **kwargs):
        """ This will be called for any version of Django < 1.6 """
        qset = super(RelatedManager, self).get_query_set(*args, **kwargs)
        if self._select:
            qset = qset.select_related(*self._select)
        if self._prefetch:
            qset = qset.prefetch_related(*self._prefetch)
        return qset

    def get_queryset(self, *args, **kwargs):
        """ This will be called for versions of Django >= 1.6 """
        qset = super(RelatedManager, self).get_queryset(*args, **kwargs)
        if self._select:
            qset = qset.select_related(*self._select)
        if self._prefetch:
            qset = qset.prefetch_related(*self._prefetch)
        return qset