Add as_num() method to ChoiceSet

This commit is contained in:
Jeremy Stretch 2025-02-18 16:42:44 -05:00
parent f1d7ad82dc
commit b697a25435

View File

@ -1,3 +1,5 @@
import enum
from django.conf import settings from django.conf import settings
from django.utils.translation import gettext_lazy as _ from django.utils.translation import gettext_lazy as _
@ -65,6 +67,23 @@ class ChoiceSet(metaclass=ChoiceSetMeta):
def values(cls): def values(cls):
return [c[0] for c in unpack_grouped_choices(cls._choices)] return [c[0] for c in unpack_grouped_choices(cls._choices)]
@classmethod
def as_enum(cls, name=None):
"""
Return the ChoiceSet as an Enum. If no name is provided, "Choices" will be stripped from the class name (if
present) and "Enum" will be appended. For example, "CircuitStatusChoices" will become "CircuitStatusEnum".
"""
name = name or f"{cls.__name__.split('Choices')[0]}Enum"
data = {}
choices = cls.values()
for attr in dir(cls):
value = getattr(cls, attr)
if attr.isupper() and value in choices:
data[attr] = value
return enum.Enum(name, data)
def unpack_grouped_choices(choices): def unpack_grouped_choices(choices):
""" """