Fixed #27910 -- Added a ChoiceEnum class for use in field choices

The class can serve as a base class for user enums, supporting
translatable human-readable names, or names automatically inferred
from the enum member name.

A ChoiceIntEnum is also probided.

Thanks Ian Foote for review.

django/db/models/fields/__init__.py

@@ -6,6 +6,8 @@
 import uuid
 import warnings
 from base64 import b64decode, b64encode
+from collections import OrderedDict
+from enum import Enum, EnumMeta
 from functools import partialmethod, total_ordering
 
 from django import forms
@@ -33,11 +35,11 @@
 
 __all__ = [
     'AutoField', 'BLANK_CHOICE_DASH', 'BigAutoField', 'BigIntegerField',
-    'BinaryField', 'BooleanField', 'CharField', 'CommaSeparatedIntegerField',
-    'DateField', 'DateTimeField', 'DecimalField', 'DurationField',
-    'EmailField', 'Empty', 'Field', 'FieldDoesNotExist', 'FilePathField',
-    'FloatField', 'GenericIPAddressField', 'IPAddressField', 'IntegerField',
-    'NOT_PROVIDED', 'NullBooleanField', 'PositiveIntegerField',
+    'BinaryField', 'BooleanField', 'CharField', 'ChoiceEnum', 'ChoiceIntEnum',
+    'CommaSeparatedIntegerField', 'DateField', 'DateTimeField', 'DecimalField',
+    'DurationField', 'EmailField', 'Empty', 'Field', 'FieldDoesNotExist',
+    'FilePathField', 'FloatField', 'GenericIPAddressField', 'IPAddressField',
+    'IntegerField', 'NOT_PROVIDED', 'NullBooleanField', 'PositiveIntegerField',
     'PositiveSmallIntegerField', 'SlugField', 'SmallIntegerField', 'TextField',
     'TimeField', 'URLField', 'UUIDField',
 ]
@@ -56,6 +58,86 @@ class NOT_PROVIDED:
 BLANK_CHOICE_DASH = [("", "---------")]
 
 
+class ChoiceEnumMeta(EnumMeta):
+    """A metaclass to support ChoiceEnum"""
+    def __new__(metacls, classname, bases, classdict):
+        choices = OrderedDict()
+        for key in classdict._member_names:
+
+            source_value = classdict[key]
+            if isinstance(source_value, (list, tuple)):
+                try:
+                    val, display = source_value
+                except ValueError:
+                    raise ValueError(f"Invalid ChoiceEnum member '{key}'")
+            else:
+                val = source_value
+                display = key.replace("_", " ").title()
+
+            choices[val] = display
+            # Use dict.__setitem__() to suppress defenses against
+            # double assignment in enum's classdict
+            dict.__setitem__(classdict, key, val)
+        cls = super().__new__(metacls, classname, bases, classdict)
+        cls._choices = choices
+        cls.choices = tuple(choices.items())
+        return cls
+
+
+class ChoiceEnum(Enum, metaclass=ChoiceEnumMeta):
+    """
+    A class suitable for using as an enum with translatable choices
+
+    The typical use is similar to the stdlib's enums, with three
+    modifications:
+    * Instead of values in the enum, we use "(value, display)" tuples.
+      The "display" can be a lazy translatable string.
+    * We add a class method "choices()" which returns a value suitable
+      for use as "choices" in a Django field definition.
+    * We add a property "display" on enum values, to return the display
+      specified.
+
+    Thus, the definition of the Enum class can look like:
+
+    class YearInSchool(ChoiceEnum):
+        FRESHMAN = 'FR', _('Freshman')
+        SOPHOMORE = 'SO', _('Sophomore')
+        JUNIOR = 'JR', _('Junior')
+        SENIOR = 'SR', _('Senior')
+
+    or even
+
+    class Suit(ChoiceIntEnum):
+        DIAMOND = 1, _('Diamond')
+        SPADE   = 2, _('Spade')
+        HEART   = 3, _('Heart')
+        CLUB    = 4, _('Club')
+
+    A field could be defined as
+
+    class Card(models.Model):
+        suit = models.IntegerField(choices=Suit.choices)
+
+    Suit.HEART, Suit['HEART'] and Suit(3) work as expected, while
+    Suit.HEART.display is a pretty, translatable string.
+    """
+    @property
+    def display(self):
+        return self._choices[self]
+
+    @classmethod
+    def validate(cls, value, message=_("Invalid choice")):
+        try:
+            cls(value)
+        except ValueError:
+            raise exceptions.ValidationError(message)
+
+
+class ChoiceIntEnum(int, ChoiceEnum):
+    """Included for symmetry with IntEnum"""
+    pass
+
+
 def _load_field(app_label, model_name, field_name):
     return apps.get_model(app_label, model_name)._meta.get_field(field_name)
 

docs/ref/models/fields.txt

@@ -96,36 +96,42 @@ and the second element is the human-readable name. For example::
         ('SR', 'Senior'),
     )
 
-Generally, it's best to define choices inside a model class, and to
-define a suitably-named constant for each value::
+Django provides a special Enum type you can subclass to define this iterable
+in a concise way::
 
     from django.db import models
+    from django.utils.translation import ugettext as _
+
+    class YearInSchool(str, models.ChoiceEnum):
+        FRESHMAN = 'FR', _('Freshman')
+        SOPHOMORE = 'SO', _('Sophomore')
+        JUNIOR = 'JR', _('Junior')
+        SENIOR = 'SR', _('Senior')
+
 
     class Student(models.Model):
-        FRESHMAN = 'FR'
-        SOPHOMORE = 'SO'
-        JUNIOR = 'JR'
-        SENIOR = 'SR'
-        YEAR_IN_SCHOOL_CHOICES = (
-            (FRESHMAN, 'Freshman'),
-            (SOPHOMORE, 'Sophomore'),
-            (JUNIOR, 'Junior'),
-            (SENIOR, 'Senior'),
-        )
         year_in_school = models.CharField(
             max_length=2,
-            choices=YEAR_IN_SCHOOL_CHOICES,
-            default=FRESHMAN,
+            choices=YearInSchool.choices,
+            default=YearInSchool.FRESHMAN,
         )
 
         def is_upperclass(self):
-            return self.year_in_school in (self.JUNIOR, self.SENIOR)
+            return self.year_in_school in (YearInSchool.JUNIOR, YearInSchool.SENIOR)
+
+If you do not need to have the human-readable names translated, you
+can have them inferred from the enum member name (replacing
+underscores to spaces and using title-case)::
 
-Though you can define a choices list outside of a model class and then
-refer to it, defining the choices and names for each choice inside the
-model class keeps all of that information with the class that uses it,
-and makes the choices easy to reference (e.g, ``Student.SOPHOMORE``
-will work anywhere that the ``Student`` model has been imported).
+    class YearInSchool(str, models.ChoiceEnum):
+        FRESHMAN = 'FR'
+        SOPHOMORE = 'SO'
+        JUNIOR = 'JR'
+        SENIOR = 'SR'
+
+Since the case where the enum values need to be integers is extremely
+common, Django provides a ``ChoiceIntEnum`` class (inheriting int and
+``ChoiceEnum``).
 
 You can also collect your available choices into named groups that can
 be used for organizational purposes::
@@ -168,6 +174,11 @@ containing ``None``; e.g. ``(None, 'Your String For Display')``.
 Alternatively, you can use an empty string instead of ``None`` where this makes
 sense - such as on a :class:`~django.db.models.CharField`.
 
+.. versionadded: 3.0
+
+   The ``ChoiceEnum`` and ``ChoiceIntEnum`` classes were added in Django 3.0.
+
+
 ``db_column``
 -------------
 

docs/releases/3.0.txt

@@ -189,6 +189,9 @@ Models
   :class:`~django.db.models.functions.Trunc` database functions determines the
   treatment of nonexistent and ambiguous datetimes.
 
+* Added the ``ChoiceEnum`` and ``ChoiceIntEnum`` helpers for defining field
+  :attr:`~django.db.models.Field.choices`.
+
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 

tests/model_fields/test_charfield.py

@@ -28,9 +28,21 @@ def test_emoji(self):
         p.refresh_from_db()
         self.assertEqual(p.title, 'Smile 😀')
 
+    def test_assignment_from_choice_enum(self):
+        class Choices(str, models.ChoiceEnum):
+            C = 'Carnival!'
+            F = 'Festival!'
+        p = Post.objects.create(title=Choices.C, body=Choices.F)
+        p.refresh_from_db()
+        self.assertEqual(p.title, 'Carnival!')
+
 
 class ValidationTests(SimpleTestCase):
 
+    class CharChoices(str, models.ChoiceEnum):
+        C = 'c', 'C'
+        D = 'd'
+
     def test_charfield_raises_error_on_empty_string(self):
         f = models.CharField()
         with self.assertRaises(ValidationError):
@@ -49,6 +61,15 @@ def test_charfield_with_choices_raises_error_on_invalid_choice(self):
         with self.assertRaises(ValidationError):
             f.clean('not a', None)
 
+    def test_charfield_with_choices_cleans_valid_choice_from_enum(self):
+        f = models.CharField(max_length=1, choices=self.CharChoices.choices)
+        self.assertEqual('c', f.clean('c', None))
+
+    def test_charfield_with_choices_raises_error_on_invalid_choice_from_enum(self):
+        f = models.CharField(max_length=1, choices=self.CharChoices.choices)
+        with self.assertRaises(ValidationError):
+            f.clean('a', None)
+
     def test_charfield_raises_error_on_empty_input(self):
         f = models.CharField(null=False)
         with self.assertRaises(ValidationError):

tests/model_fields/test_choiceenum.py

@@ -0,0 +1,79 @@
+from django.core.exceptions import ValidationError
+from django.db import models
+from django.test import SimpleTestCase
+
+
+class TestChoiceEnumDefinition(SimpleTestCase):
+
+    def test_simple_definition_int(self):
+
+        class MyEnum(models.ChoiceIntEnum):
+            ONE = 1, 'One'
+            TWO = 2, 'Two'
+
+        self.assertEqual(MyEnum.ONE, 1)
+        self.assertEqual(MyEnum.TWO.display, 'Two')
+        self.assertSequenceEqual(MyEnum.choices, [(1, 'One'), (2, 'Two')])
+
+    def test_simple_definition_int_default_display(self):
+        class MyEnum(models.ChoiceIntEnum):
+            ONE = 1, 'One'
+            ONE_PLUS_ONE = 2
+            THREE = 3
+
+        self.assertEqual(MyEnum.ONE.display, 'One')
+        self.assertEqual(MyEnum.ONE_PLUS_ONE.display, 'One Plus One')
+        self.assertSequenceEqual(MyEnum.choices, [(1, 'One'), (2, 'One Plus One'), (3, 'Three')])
+
+    def test_simple_definition_str(self):
+
+        class MyEnum(str, models.ChoiceEnum):
+            ONE = '1', 'One'
+            TWO = '2', 'Two'
+
+        self.assertEqual(MyEnum.ONE, '1')
+        self.assertEqual(MyEnum.TWO.display, 'Two')
+        self.assertSequenceEqual(MyEnum.choices, [('1', 'One'), ('2', 'Two')])
+
+    def test_simple_definition_str_default_display(self):
+        class MyEnum(str, models.ChoiceEnum):
+            ONE = '1', 'One'
+            ONE_PLUS_ONE = '2'
+            THREE = '3'
+
+        self.assertEqual(MyEnum.ONE.display, 'One')
+        self.assertEqual(MyEnum.ONE_PLUS_ONE.display, 'One Plus One')
+        self.assertSequenceEqual(MyEnum.choices, [('1', 'One'), ('2', 'One Plus One'), ('3', 'Three')])
+
+    def test_invalid_definition(self):
+        with self.assertRaisesRegex(ValueError, '.*[Ii]nvalid.*ONE.*'):
+            class MyEnum(models.ChoiceIntEnum):
+                ONE = 1, 'One', 'One too many'
+
+
+class TestChoiceEnumValidation(SimpleTestCase):
+
+    class IntEnum(models.ChoiceIntEnum):
+        ONE = 1, 'One'
+        TWO = 2, 'Two'
+
+    class StrEnum(str, models.ChoiceEnum):
+        ONE = '1', 'One'
+        ONE_PLUS_ONE = '2'
+        THREE = '3'
+
+    def test_validates_valid_value(self):
+        self.IntEnum.validate(1)
+        self.StrEnum.validate('2')
+
+    def test_invalidates_correctly_typed_value(self):
+        with self.assertRaises(ValidationError):
+            self.IntEnum.validate(7)
+        with self.assertRaises(ValidationError):
+            self.StrEnum.validate('monster')
+
+    def test_invalidates_incorrectly_typed_value(self):
+        with self.assertRaises(ValidationError):
+            self.IntEnum.validate('not int')
+        with self.assertRaises(ValidationError):
+            self.StrEnum.validate(b'2')

tests/model_fields/test_integerfield.py

@@ -163,6 +163,9 @@ def test_negative_values(self):
 
 class ValidationTests(SimpleTestCase):
 
+    class Choices(models.ChoiceIntEnum):
+        A = 1
+
     def test_integerfield_cleans_valid_string(self):
         f = models.IntegerField()
         self.assertEqual(f.clean('2', None), 2)
@@ -196,3 +199,14 @@ def test_integerfield_validates_zero_against_choices(self):
         f = models.IntegerField(choices=((1, 1),))
         with self.assertRaises(ValidationError):
             f.clean('0', None)
+
+    def test_integerfield_with_choices_cleans_valid_choice_from_enum(self):
+        f = models.IntegerField(choices=self.Choices.choices)
+        self.assertEqual(1, f.clean('1', None))
+
+    def test_charfield_with_choices_raises_error_on_invalid_choice_from_enum(self):
+        f = models.IntegerField(choices=self.Choices.choices)
+        with self.assertRaises(ValidationError):
+            f.clean('A', None)
+        with self.assertRaises(ValidationError):
+            f.clean('3', None)