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
  • Star 3 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • 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)
"""
@nitrag
Copy link

nitrag commented Jun 16, 2018

Do you not get this error in Admin when editing?

  File "/Users/ryan/Documents/asdf/django/lib/python3.6/site-packages/django/forms/widgets.py", line 677, in get_context
    context = super().get_context(name, value, attrs)
  File "/Users/ryan/Documents/asdf/django/lib/python3.6/site-packages/django/forms/widgets.py", line 637, in get_context
    context['widget']['optgroups'] = self.optgroups(name, context['widget']['value'], attrs)
  File "/Users/ryan/Documents/asdf/django/lib/python3.6/site-packages/django/forms/widgets.py", line 585, in optgroups
    for index, (option_value, option_label) in enumerate(self.choices):
  File "/Users/ryan/Documents/asdf/django/lib/python3.6/site-packages/django/forms/fields.py", line 747, in __iter__
    yield from self.choices_func()
TypeError: __call__() missing 1 required positional argument: 'value'

@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