CODEWITHSUNDEEP
pythondjangodjango-models
Johnny Metz
Johnny Metz

Reputation: 5955

Sort Django queryset with custom order

Let's say I have the following Django model:

class Todo(models.Model):
    class Priority(models.IntegerChoices):
        HIGH = 1, "High"
        LOW = 2, "Low"
        MEDIUM = 3, "Medium"

    title = models.CharField(max_length=255)
    priority = models.PositiveSmallIntegerField(choices=Priority.choices, db_index=True)

And the following test data via factory-boy:

import factory
from factory.django import DjangoModelFactory

class TodoFactory(DjangoModelFactory):
    class Meta:
        model = Todo

    title = factory.Sequence(lambda n: f"Todo {n}")

def test_custom_db_sort():
    todo1 = TodoFactory(priority=Todo.Priority.MEDIUM)
    todo2 = TodoFactory(priority=Todo.Priority.HIGH)
    todo3 = TodoFactory(priority=Todo.Priority.LOW)
    todo4 = TodoFactory(priority=Todo.Priority.MEDIUM)
    todo5 = TodoFactory(priority=Todo.Priority.HIGH)

I want to fetch the Todo objects from the database and sort from HIGH -> LOW priority. Unfortunately Todo.Priority isn't ordered correctly so I can't sort on that. I've come up with the following solution:

PREFERENCE = {
    Priority.HIGH: 1,
    Priority.MEDIUM: 2,
    Priority.LOW: 3,
}
result = sorted(
    Todo.objects.all(),
    key=lambda x: [PREFERENCE[x.priority], x.id]
)

This works but I'd rather sort in the database for several reasons (e.g. faster, better to do work in the DB rather than Python). Is there anyway to perform this query in the database?

Upvotes: 0

Views: 1990

Answers (1)

JPG
JPG

Reputation: 88429

Use Conditional expression (Case(...), When(...)) expressions

from django.db.models import Case, When, Value

new_priority = Case(
    When(priority=Todo.Priority.HIGH, then=Value(1)),
    When(priority=Todo.Priority.MEDIUM, then=Value(2)),
    When(priority=Todo.Priority.LOW, then=Value(3)),
)

qs = Todo.objects.annotate(new_priority=new_priority).order_by("new_priority")

If I have a choice, I would rather update the Priority choice values instead of annotating the expression.

Upvotes: 2

Related Questions