enums.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import enum
  2. from django.utils.functional import Promise
  3. __all__ = ['Choices', 'IntegerChoices', 'TextChoices']
  4. class ChoicesMeta(enum.EnumMeta):
  5. """A metaclass for creating a enum choices."""
  6. def __new__(metacls, classname, bases, classdict):
  7. labels = []
  8. for key in classdict._member_names:
  9. value = classdict[key]
  10. if (
  11. isinstance(value, (list, tuple)) and
  12. len(value) > 1 and
  13. isinstance(value[-1], (Promise, str))
  14. ):
  15. *value, label = value
  16. value = tuple(value)
  17. else:
  18. label = key.replace('_', ' ').title()
  19. labels.append(label)
  20. # Use dict.__setitem__() to suppress defenses against double
  21. # assignment in enum's classdict.
  22. dict.__setitem__(classdict, key, value)
  23. cls = super().__new__(metacls, classname, bases, classdict)
  24. cls._value2label_map_ = dict(zip(cls._value2member_map_, labels))
  25. # Add a label property to instances of enum which uses the enum member
  26. # that is passed in as "self" as the value to use when looking up the
  27. # label in the choices.
  28. cls.label = property(lambda self: cls._value2label_map_.get(self.value))
  29. cls.do_not_call_in_templates = True
  30. return enum.unique(cls)
  31. def __contains__(cls, member):
  32. if not isinstance(member, enum.Enum):
  33. # Allow non-enums to match against member values.
  34. return member in {x.value for x in cls}
  35. return super().__contains__(member)
  36. @property
  37. def names(cls):
  38. empty = ['__empty__'] if hasattr(cls, '__empty__') else []
  39. return empty + [member.name for member in cls]
  40. @property
  41. def choices(cls):
  42. empty = [(None, cls.__empty__)] if hasattr(cls, '__empty__') else []
  43. return empty + [(member.value, member.label) for member in cls]
  44. @property
  45. def labels(cls):
  46. return [label for _, label in cls.choices]
  47. @property
  48. def values(cls):
  49. return [value for value, _ in cls.choices]
  50. class Choices(enum.Enum, metaclass=ChoicesMeta):
  51. """Class for creating enumerated choices."""
  52. def __str__(self):
  53. """
  54. Use value when cast to str, so that Choices set as model instance
  55. attributes are rendered as expected in templates and similar contexts.
  56. """
  57. return str(self.value)
  58. class IntegerChoices(int, Choices):
  59. """Class for creating enumerated integer choices."""
  60. pass
  61. class TextChoices(str, Choices):
  62. """Class for creating enumerated string choices."""
  63. def _generate_next_value_(name, start, count, last_values):
  64. return name