enums.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123
  1. import enum
  2. import warnings
  3. from django.utils.deprecation import RemovedInDjango60Warning
  4. from django.utils.functional import Promise
  5. from django.utils.version import PY311, PY312
  6. if PY311:
  7. from enum import EnumType, IntEnum, StrEnum
  8. from enum import property as enum_property
  9. else:
  10. from enum import EnumMeta as EnumType
  11. from types import DynamicClassAttribute as enum_property
  12. class ReprEnum(enum.Enum):
  13. def __str__(self):
  14. return str(self.value)
  15. class IntEnum(int, ReprEnum):
  16. pass
  17. class StrEnum(str, ReprEnum):
  18. pass
  19. __all__ = ["Choices", "IntegerChoices", "TextChoices"]
  20. class ChoicesType(EnumType):
  21. """A metaclass for creating a enum choices."""
  22. def __new__(metacls, classname, bases, classdict, **kwds):
  23. labels = []
  24. for key in classdict._member_names:
  25. value = classdict[key]
  26. if (
  27. isinstance(value, (list, tuple))
  28. and len(value) > 1
  29. and isinstance(value[-1], (Promise, str))
  30. ):
  31. *value, label = value
  32. value = tuple(value)
  33. else:
  34. label = key.replace("_", " ").title()
  35. labels.append(label)
  36. # Use dict.__setitem__() to suppress defenses against double
  37. # assignment in enum's classdict.
  38. dict.__setitem__(classdict, key, value)
  39. cls = super().__new__(metacls, classname, bases, classdict, **kwds)
  40. for member, label in zip(cls.__members__.values(), labels):
  41. member._label_ = label
  42. return enum.unique(cls)
  43. if not PY312:
  44. def __contains__(cls, member):
  45. if not isinstance(member, enum.Enum):
  46. # Allow non-enums to match against member values.
  47. return any(x.value == member for x in cls)
  48. return super().__contains__(member)
  49. @property
  50. def names(cls):
  51. empty = ["__empty__"] if hasattr(cls, "__empty__") else []
  52. return empty + [member.name for member in cls]
  53. @property
  54. def choices(cls):
  55. empty = [(None, cls.__empty__)] if hasattr(cls, "__empty__") else []
  56. return empty + [(member.value, member.label) for member in cls]
  57. @property
  58. def labels(cls):
  59. return [label for _, label in cls.choices]
  60. @property
  61. def values(cls):
  62. return [value for value, _ in cls.choices]
  63. class Choices(enum.Enum, metaclass=ChoicesType):
  64. """Class for creating enumerated choices."""
  65. if PY311:
  66. do_not_call_in_templates = enum.nonmember(True)
  67. else:
  68. @property
  69. def do_not_call_in_templates(self):
  70. return True
  71. @enum_property
  72. def label(self):
  73. return self._label_
  74. # A similar format was proposed for Python 3.10.
  75. def __repr__(self):
  76. return f"{self.__class__.__qualname__}.{self._name_}"
  77. class IntegerChoices(Choices, IntEnum):
  78. """Class for creating enumerated integer choices."""
  79. pass
  80. class TextChoices(Choices, StrEnum):
  81. """Class for creating enumerated string choices."""
  82. @staticmethod
  83. def _generate_next_value_(name, start, count, last_values):
  84. return name
  85. def __getattr__(name):
  86. if name == "ChoicesMeta":
  87. warnings.warn(
  88. "ChoicesMeta is deprecated in favor of ChoicesType.",
  89. RemovedInDjango60Warning,
  90. stacklevel=2,
  91. )
  92. return ChoicesType
  93. raise AttributeError(f"module {__name__!r} has no attribute {name!r}")