In Django, can you add a method to querysets?

Question

In Django, if I have a model class, e.g.

from django.db import models

class Transaction(models.Model):
    ...

then if I want to add methods to the model, to store e.g. reasonably complex filters, I can add a custom model manager, e.g.

class TransactionManager(models.Manager):

    def reasonably_complex_filter(self):
        return self.get_query_set().filter(...)


class Transaction(models.Model):
    objects = TransactionManager()

And then I can do:

>>> Transaction.objects.reasonably_complex_filter()

Is there any way I can add a custom method that can be chained to the end of a query set from the model?

i.e. add the custom method in such a way that I can do this:

>>> Transaction.objects.filter(...).reasonably_complex_filter()

Asked By: Paul D. Waite

||

Source

Answer 1

You need to add methods to the QuerySet which you eventually end up with. So you need to create and use a QuerySet subclass which has the methods you define wherever you want this functionality.

I found this tutorial which explains how to do it and reasons why you might want to:

https://web.archive.org/web/20160329131857/http://adam.gomaa.us/blog/2009/feb/16/subclassing-django-querysets/index.html

Answered By: Marcus Whybrow

Answer 2

You could modify the get_query_set() method to return a custom QuerySet, adding the methods you require. In your case, you would use:

class TransactionManager(models.Manager):
    def get_query_set(self):
        return TransactionQuerySet(self.model)

class TransactionQuerySet(models.query.QuerySet):
    def reasonably_complex_filter(self):
        return self.filter(...)

I’ve seen examples subclassing the TransactionQuerySet into the Transaction model, or in the related Manager, but that is entirely up to you.

edit: I seem to have overlooked the fact that objects first references to the TransactionManager and therefore Transaction.objects.reasonably_complex_filter() isn’t possible in my implementation. This can be fixed in three ways:

Implement the reasonably_complex_filter in both the Manager and the QuerySet;
Use Transaction.objects.all().reasonably_complex_filter() when that’s the only filter required;
Refer to Marcus Whybrow’s answer for a solution that will implement the method in both the QuerySet and the Manager without code duplication.

It depends on the application which option is the most desireable, although I would strongly recommend against code duplication (although you could use a global method to overcome this). Though, the last option could be too costly in terms of overhead if you only require this kind of practice once, or if you only intend to use the complex filter in combination with another filter.

Answered By: ralphje

Answer 3

I actually ended up going with another method. It turned out I only needed to chain filter calls onto the end of my custom method, so I amended my method to take arbitrary keyword arguments, and pass them through to a filter() call on the end of my reasonably complex query:

class TransactionManager(models.Manager):

    def reasonably_complex_filter(self, **kwargs):
        return self.get_query_set().filter(...).filter(**kwargs)

Seems to work fine for my purposes, and is a bit simpler than subclassing QuerySet.

Answered By: Paul D. Waite

Answer 4

This is a complete solution that is known to work in Django 1.3, courtesy of Zach Smith and Ben.

class Entry(models.Model):
    objects = EntryManager() # don't forget this

    is_public = models.BooleanField()
    owner = models.ForeignKey(User)


class EntryManager(models.Manager):
    '''Use this class to define methods just on Entry.objects.'''
    def get_query_set(self):
        return EntryQuerySet(self.model)

    def __getattr__(self, name, *args):
        if name.startswith("_"): 
            raise AttributeError
        return getattr(self.get_query_set(), name, *args) 

    def get_stats(self):
        '''A sample custom Manager method.'''
        return { 'public_count': self.get_query_set().public().count() }


class EntryQuerySet(models.query.QuerySet):
    '''Use this class to define methods on queryset itself.'''
    def public(self):
        return self.filter(is_public=True)

    def by(self, owner):
        return self.filter(owner=owner)


stats = Entry.objects.get_stats()    
my_entries = Entry.objects.by(request.user).public()

Note: the get_query_set() method is now deprecated in Django 1.6; get_queryset() should be used instead in this case.

Answered By: Dan Abramov

Answer 5

As of django 1.7, the ability to use a query set as a manager was added:

class PersonQuerySet(models.QuerySet):
    def authors(self):
        return self.filter(role='A')

    def editors(self):
        return self.filter(role='E')

class Person(models.Model):
    first_name = models.CharField(max_length=50)
    last_name = models.CharField(max_length=50)
    role = models.CharField(max_length=1, choices=(('A', _('Author')),
                                                   ('E', _('Editor'))))
    people = PersonQuerySet.as_manager()

Resulting the following:

Person.people.authors(last_name='Dahl')

In addition, the ability to add custom lookups was also added.

Answered By: Burhan Khalid

Answer 6

If you need both custom Manager methods and custom QuerySet methods, you can use from_queryset.

class BaseManager(models.Manager):
    def manager_only_method(self):
        return

class CustomQuerySet(models.QuerySet):
    def manager_and_queryset_method(self):
        return

class MyModel(models.Model):
    objects = BaseManager.from_queryset(CustomQuerySet)()

https://docs.djangoproject.com/en/2.1/topics/db/managers/#from-queryset

Answered By: Stian Jensen

Answer 7

Simplest version for Django 2.0+

class TransactionQuerySet(models.QuerySet):
    def reasonably_complex_filter(self):
        return self.filter(...)

class Transaction(models.Model):
    ...

    objects = models.Manager.from_queryset(TransactionQuerySet)()

Answered By: Mark Mishyn

In Django, can you add a method to querysets?

Question:

Answers: