blog: Subclassing Django QuerySets

James Bennett wrote a great explanation last year of when it's a good idea to use class methods versus custom managers in Django models. I'd like to showcase a snippet from Django Snippets, which, combined with one extra class, allows you to add expressive, chainable querysets to your managers.

Why Not Class Methods?

I spent about the first year of my Django use avoiding managers; I was happy with class methods, mostly because those were what I was already familiar with.

Ultimately, James' reasons are correct; the examples I'm about to show work best when they're done as model managers. Trying to make class methods out of them would require the use of mixin classes - not a great idea for the general case of adding query functionality, and more verbose to boot.

And Why Not Plain Custom Managers?

Django already allows you to define custom managers. The problem is that once you're past the first call after the manager - Foo.objects.all() - you're dealing with a plain QuerySet, which means you lose all your custom methods. So if you wanted to do MyModel.objects.popular().public() - well, you can't. After that first .popular() call, you're left with just the normal QuerySet methods - .filter(), .exclude(), etc.

One thing you can do is to override .get_query_set() to return a queryset that's an instance of a QuerySet subclass, which you can then add your own methods to. The problem with that is that you now have to define your methods twice - once on the Manager, and once on the QuerySet. Alternately, you can just remember to always put at least one .all() call in before you start using your custom QuerySet methods... but we can avoid that, too.

Let's get to it.

The Code

The first part of this comes from a comment on Snippet #562. This is a piece of support code, which you'd want to keep in a utility module of some sort.

from django.db import models
# http://www.djangosnippets.org/snippets/562/#c673
class QuerySetManager(models.Manager):
    # http://docs.djangoproject.com/en/dev/topics/db/managers/#using-managers-for-related-object-access
    # Not working cause of:
    # http://code.djangoproject.com/ticket/9643
    use_for_related_fields = True
    def __init__(self, qs_class=models.query.QuerySet):
        self.queryset_class = qs_class
        super(QuerySetManager, self).__init__()

    def get_query_set(self):
        return self.queryset_class(self.model)

    def __getattr__(self, attr, *args):
        try:
            return getattr(self.__class__, attr, *args)
        except AttributeError:
            return getattr(self.get_query_set(), attr, *args)

This allows you to define django.db.models.query.QuerySet subclasses, pass them to the __init__ of this manager, and have the manager proxy unknown attributes through to the queryset. This is essentially what django.db.models.manager.Manager does - albeit, by explicitly defining the methods that are proxied through to .get_query_set(), not with a __getattr__ hook. That's the right choice for Django core - but my standards are considerably lower.

One problem with what we have so far - and it's not a big one - is that, at this point, you're still importing this QuerySetManager into each models file that you intend to use it:

# foobar/models.py
from django.db import models
from myproject.utils import QuerySetManager


class CustomQuerySet(models.query.QuerySet):
    def some_method(self):
        return self.filter(field__startswith="foo")


class MyModel(models.Model):
    field = models.CharField(max_length=127)

    objects = QuerySetManager(CustomQuerySet)

This is also pretty easy to fix, with another support class:

from django.db import models
class QuerySet(models.query.QuerySet):
    """Base QuerySet class for adding custom methods that are made
    available on both the manager and subsequent cloned QuerySets"""

    @classmethod
    def as_manager(cls, ManagerClass=QuerySetManager):
        return ManagerClass(cls)

The key here is as_manager, which lets us change the above models.py example like so:

# foobar/models.py
from django.db import models
from myproject.utils import QuerySet


class CustomQuerySet(QuerySet):
    def some_method(self):
        return self.filter(field__startswith="foo")


class MyModel(models.Model):
    field = models.CharField(max_length=127)

    objects = CustomQuerySet.as_manager()

And at this point, we have all the support code that should be needed. Personally, I like to add some additional methods to the support QuerySet class:

def random(self):
    return self.order_by("?")

def first(self):
    return self[0]

but that's up to you.