Django model field choices with an inner class enum

2018-05-02 (Wednesday) | 300 words (~2 minutes reading)

It’s common to define the possible choices for a character field on Django models like this (from the docs):

from django.db import models

class Student(models.Model):
    FRESHMAN = 'FR'
    SOPHOMORE = 'SO'
    JUNIOR = 'JR'
    SENIOR = 'SR'
    YEAR_IN_SCHOOL_CHOICES = (
        (FRESHMAN, 'Freshman'),
        (SOPHOMORE, 'Sophomore'),
        (JUNIOR, 'Junior'),
        (SENIOR, 'Senior'),
    )
    year_in_school = models.CharField(
        max_length=2,
        choices=YEAR_IN_SCHOOL_CHOICES,
        default=FRESHMAN,
    )

This works fine and it makes constants available for the values, so that elsewhere you can refer to Student.FRESHMAN. It is quite verbose, though, and doesn’t group the constants together very well.

One way to group the constants more clearly is to use an inner class:

from django.db import models

class Student(models.Model):
    class Years:
        FRESHMAN = 'FR'
        SOPHOMORE = 'SO'
        JUNIOR = 'JR'
        SENIOR = 'SR'
        CHOICES = (
            (FRESHMAN, 'Freshman'),
            (SOPHOMORE, 'Sophomore'),
            (JUNIOR, 'Junior'),
            (SENIOR, 'Senior'),
        )
    year_in_school = models.CharField(
        max_length=2,
        choices=Years.CHOICES,
        default=FRESHMAN,
    )

Now the constants are clearly part of one group. We can also refer to them intuitively elsewhere with Student.Years.FRESHMAN. It’s still a little verbose, though, as we have to manually re-define the constants as a choices tuple.

We can make one further improvement on this by using an enum for the inner class:

from enum import Enum
from django.db import models

class ChoiceEnum(Enum):
    @classmethod
    def choices(cls):
        return [(choice.name, choice.value) for choice in cls]

class Student(models.Model):
    class Years(ChoiceEnum):
        FR = 'Freshman'
        SO = 'Sophomore'
        JR = 'Junior'
        SR = 'Senior'

    year_in_school = models.CharField(
        max_length=2,
        choices=Years.choices(),
        default=FRESHMAN,
    )

Now we can easily refer to Student.Years.FR.value elsewhere, and the Student model is kept simple and readable.

It also has the nice consequence that we can use the enum as a type-hint elsewhere:

from some_module import Student

def find_students_by_year(year: Student.Year):
    pass

References


Tech mentioned