pythondjangodjango-modelsdjango-querysetdjango-managers

Custom QuerySet and Manager without breaking DRY?


I'm trying to find a way to implement both a custom QuerySet and a custom Manager without breaking DRY. This is what I have so far:

class MyInquiryManager(models.Manager):
    def for_user(self, user):
        return self.get_query_set().filter(
                    Q(assigned_to_user=user) |
                    Q(assigned_to_group__in=user.groups.all())
                )

class Inquiry(models.Model):   
    ts = models.DateTimeField(auto_now_add=True)
    status = models.ForeignKey(InquiryStatus)
    assigned_to_user = models.ForeignKey(User, blank=True, null=True)
    assigned_to_group = models.ForeignKey(Group, blank=True, null=True)
    objects = MyInquiryManager()

This works fine, until I do something like this:

inquiries = Inquiry.objects.filter(status=some_status)
my_inquiry_count = inquiries.for_user(request.user).count()

This promptly breaks everything because the QuerySet doesn't have the same methods as the Manager. I've tried creating a custom QuerySet class, and implementing it in MyInquiryManager, but I end up replicating all of my method definitions.

I also found this snippet which works, but I need to pass in the extra argument to for_user so it breaks down because it relies heavily on redefining get_query_set.

Is there a way to do this without redefining all of my methods in both the QuerySet and the Manager subclasses?


Solution

  • Django has changed! Before using the code in this answer, which was written in 2009, be sure to check out the rest of the answers and the Django documentation to see if there is a more appropriate solution.


    The way I've implemented this is by adding the actual get_active_for_account as a method of a custom QuerySet. Then, to make it work off the manager, you can simply trap the __getattr__ and return it accordingly

    To make this pattern re-usable, I've extracted out the Manager bits to a separate model manager:

    custom_queryset/models.py

    from django.db import models
    from django.db.models.query import QuerySet
    
    class CustomQuerySetManager(models.Manager):
        """A re-usable Manager to access a custom QuerySet"""
        def __getattr__(self, attr, *args):
            try:
                return getattr(self.__class__, attr, *args)
            except AttributeError:
                # don't delegate internal methods to the queryset
                if attr.startswith('__') and attr.endswith('__'):
                    raise
                return getattr(self.get_query_set(), attr, *args)
    
        def get_query_set(self):
            return self.model.QuerySet(self.model, using=self._db)
    

    Once you've got that, on your models all you need to do is define a QuerySet as a custom inner class and set the manager to your custom manager:

    your_app/models.py

    from custom_queryset.models import CustomQuerySetManager
    from django.db.models.query import QuerySet
    
    class Inquiry(models.Model):
        objects = CustomQuerySetManager()
    
        class QuerySet(QuerySet):
            def active_for_account(self, account, *args, **kwargs):
                return self.filter(account=account, deleted=False, *args, **kwargs)
    

    With this pattern, any of these will work:

    >>> Inquiry.objects.active_for_account(user)
    >>> Inquiry.objects.all().active_for_account(user)
    >>> Inquiry.objects.filter(first_name='John').active_for_account(user)
    

    UPD if you are using it with custom user(AbstractUser), you need to change
    from

    class CustomQuerySetManager(models.Manager):
    

    to

    from django.contrib.auth.models import UserManager
    
    class CustomQuerySetManager(UserManager):
        ***