Skip to content

Instantly share code, notes, and snippets.

@tadeo
Forked from treyhunner/choice_enum.py
Last active August 15, 2018 23:51
Show Gist options
  • Save tadeo/0d3e94380353bb58451e0b5b72a9588f to your computer and use it in GitHub Desktop.
Save tadeo/0d3e94380353bb58451e0b5b72a9588f to your computer and use it in GitHub Desktop.
Enum for use with Django ChoiceField
from enum import Enum, EnumMeta
class ChoicesEnumMeta(EnumMeta):
def __iter__(self):
return ((choice.name, choice.value) for choice in super().__iter__())
class ChoicesEnum(Enum, metaclass=ChoicesEnumMeta):
def __str__(self):
return self.name
"""
Enum for Django ChoiceField use.
Usage:
class Colors(ChoicesEnum):
red = "Red"
green = "Green"
blue = "Blue"
class Car(models.Model):
color = models.CharField(max_length=20, choices=Colors)
red_cars = Car.objects.filter(color=Colors.RED)
"""
@tadeo
Copy link
Author

tadeo commented Jun 17, 2018

@nitrag I've just tried and I'm not having any issue. In my current implementation I nested the ChoicesEnum class inside the model class:

class Order(models.Model):

    class Status(ChoicesEnum):
        DRAFT = 'Draft'
        CONFIRMED = 'Confirmed'
        PROCESSED = 'Processed'
        CANCELED = 'Canceled'

    status = models.CharField(max_length=16, choices=Status, default=Status.DRAFT, verbose_name='status')

After that I'm being able of doing:

Order.objects.filter(status=Order.Status.DRAFT).count()

The only two issues I faced were:

  1. I have to manually edit the migration and fix the default from oakfrog.orders.models.Status('Draft') to oakfrog.orders.models.Order.Status('Draft') (It seems that Django are not handling well the nested Class when doing the migration).
  2. When testing a DRF serializer I needed to str() the option to make it assert: self.assertEqual(data['status'], str(Order.Status.CONFIRMED)). I suspect it's due DRAFT in the Enum is a symbolic name and the serializer is casting it to a plain string.

I hope it helps!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment