Skip to content

Instantly share code, notes, and snippets.

@PaperNick
Forked from treyhunner/choice_enum.py
Last active April 8, 2019 13:44
Show Gist options
  • Save PaperNick/8b6eac5301b41e35bb1cacb9efd978c4 to your computer and use it in GitHub Desktop.
Save PaperNick/8b6eac5301b41e35bb1cacb9efd978c4 to your computer and use it in GitHub Desktop.
Enum for use with Django ChoiceField
from enum import Enum, EnumMeta, IntEnum, unique
def not_callable_in_template(cls):
"""Disable the callable behavior of Enum in django templates"""
cls.do_not_call_in_templates = True
return cls
class ChoiceEnumMeta(EnumMeta):
def __iter__(cls):
return ((str(enum_member), enum_member.value)
for enum_member in super().__iter__())
@not_callable_in_template
class ChoiceEnum(Enum, metaclass=ChoiceEnumMeta):
def __str__(self):
return self.name
def __eq__(self, other):
return str(self) == str(other)
@property
def label(self):
return self.value
class IntChoiceEnumMeta(EnumMeta):
def __iter__(cls):
return ((enum_member.value, enum_member.label) for enum_member in super().__iter__())
@not_callable_in_template
@unique
class IntChoiceEnum(IntEnum, metaclass=IntChoiceEnumMeta):
"""Attributes:
__choices_label_map__ should return Dict[EnumMember, Label] where:
EnumMember - ColorEnum.RED
Label - str or lazy(str)
"""
__choices_label_map__ = {}
@property
def label(self):
return self.__choices_label_map__.get(self, '')

Usage of ChoiceEnum

Enum definition

from django.utils.translation import ugettext_lazy as _

from django_choices_enum.enums import ChoiceEnum


class ColorEnum(ChoiceEnum):
    RED = _('Red')
    BLUE = _('Blue')
    GREEN = _('Green')

In models:

from django.db import models
from django.utils.translation import ugettext_lazy as _
from .enums import ColorEnum

class ColorPalette(models.Model):
    color = models.CharField(_('Color'), max_length=30, choices=ColorEnum)

Querying:

ColorPalette.objects.filter(color=ColorEnum.RED)

In templates:

View:

def my_page(request):
    context = {
        'ColorEnum': ColorEnum,
        'color_palette': ColorPalette.objects.filter(color=ColorEnum.BLUE).first(),
    }

    return render(request, 'my_app/index.html', context)

Template:

<p>The color is {% if color_palette.color == ColorEnum.BLUE %}blue{% else %}not blue{% endif %}</p>

Will result in:

<p>The color is blue</p>

Comparing:

In [1]: ColorEnum.GREEN == 'GREEN'
Out[1]: True

Getting the label:

In [1]: ColorEnum.RED.label
Out[1]: 'Red'

Usage of IntChoiceEnum

Enum definition

from django.utils.translation import ugettext_lazy as _

from django_choices_enum.enums import IntChoiceEnum


class ShakeEnum(IntChoiceEnum):
    CHOCO = 1
    VANILLA = 2
    MINT = 3

    __choices_label_map__ = {
        CHOCO: _('Chocolatey'),
        VANILLA: _('Vanilla'),
        MINT: _('Minty'),
    }

In models:

from django.db import models
from django.utils.translation import ugettext_lazy as _
from .enums import ShakeEnum

class Smoothie(models.Model):
    shake = models.SmallIntegerField(_('Shake'), choices=ShakeEnum)

Querying:

Smoothie.objects.filter(shake=ShakeEnum.VANILLA)

In templates:

View:

def my_page(request):
    context = {
        'ShakeEnum': ShakeEnum,
        'smoothie': Smoothie.objects.filter(shake=ShakeEnum.MINT).first(),
    }

    return render(request, 'my_app/index.html', context)

Template:

<p>Your shake contains {% if smoothie.shake == ShakeEnum.MINT %}fresh mint{% else %}something not minty{% endif %}</p>

Will result in:

<p>Your shake contains fresh mint</p>

Comparing:

In [1]: ShakeEnum.MINT == 3
Out[1]: True

Getting the label:

In [1]: ShakeEnum.MINT.label
Out[1]: 'Minty'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment