Enum Named Groups Pretty

AUTHORS

@@ -908,6 +908,7 @@ answer newbie questions, and generally made Django that much better:
     Stéphane Raimbault <[email protected]>
     Stephan Jaekel <[email protected]>
     Stephen Burrows <[email protected]>
+    Steven H. Johnson <[email protected]>
     Steven L. Smith (fvox13) <[email protected]>
     Steven Noorbergen (Xaroth) <[email protected]>
     Stuart Langridge <https://www.kryogenix.org/>

django/db/models/enums.py

@@ -11,8 +11,19 @@ class ChoicesMeta(enum.EnumMeta):
 
     def __new__(metacls, classname, bases, classdict, **kwds):
         labels = []
-        for key in classdict._member_names:
+        named_groups = []
+        key_order = []
+        for key in list(classdict._member_names):
             value = classdict[key]
+            # Check if value is a named group
+            if hasattr(value, "choices"):
+                named_group = value.__label__ if hasattr(value, "__label__") else key
+                for member in value:
+                    key_order.append(member.name)
+                    classdict[member.name] = member
+                    labels.append(member.label)
+                    named_groups.append(named_group)
+                continue
             if (
                 isinstance(value, (list, tuple))
                 and len(value) > 1
@@ -22,13 +33,19 @@ def __new__(metacls, classname, bases, classdict, **kwds):
                 value = tuple(value)
             else:
                 label = key.replace("_", " ").title()
+            key_order.append(key)
             labels.append(label)
+            named_groups.append(None)
             # Use dict.__setitem__() to suppress defenses against double
             # assignment in enum's classdict.
             dict.__setitem__(classdict, key, value)
+        classdict._member_names.clear()
+        classdict._member_names.extend(key_order)
         cls = super().__new__(metacls, classname, bases, classdict, **kwds)
-        for member, label in zip(cls.__members__.values(), labels):
+        values = (cls.__members__[x] for x in key_order)
+        for member, label, named_group in zip(values, labels, named_groups):
             member._label_ = label
+            member._named_group_ = named_group
         return enum.unique(cls)
 
     def __contains__(cls, member):
@@ -44,16 +61,34 @@ def names(cls):
 
     @property
     def choices(cls):
+        choices = []
+        choice_list = choices
+        last_named_group = None
+        for member in cls:
+            if member.named_group != last_named_group:
+                last_named_group = member.named_group
+                if member.named_group:
+                    choice_list = []
+                    choices.append((member.named_group, choice_list))
+                else:
+                    choice_list = choices  # Add to toplevel
+            choice_list.append((member.value, member.label))
+        empty = [(None, cls.__empty__)] if hasattr(cls, "__empty__") else []
+        return empty + choices
+
+    @property
+    def flatchoices(cls):
+        """Flattened version of choices tuple."""
         empty = [(None, cls.__empty__)] if hasattr(cls, "__empty__") else []
         return empty + [(member.value, member.label) for member in cls]
 
     @property
     def labels(cls):
-        return [label for _, label in cls.choices]
+        return [label for _, label in cls.flatchoices]
 
     @property
     def values(cls):
-        return [value for value, _ in cls.choices]
+        return [value for value, _ in cls.flatchoices]
 
 
 class Choices(enum.Enum, metaclass=ChoicesMeta):
@@ -63,6 +98,10 @@ class Choices(enum.Enum, metaclass=ChoicesMeta):
     def label(self):
         return self._label_
 
+    @DynamicClassAttribute
+    def named_group(self):
+        return self._named_group_
+
     @property
     def do_not_call_in_templates(self):
         return True

docs/ref/models/fields.txt

@@ -229,9 +229,10 @@ modifications:
   <field-choices-enum-auto-label>` from the member name.
 * A ``.label`` property is added on values, to return the human-readable name.
 * A number of custom properties are added to the enumeration classes --
-  ``.choices``, ``.labels``, ``.values``, and ``.names`` -- to make it easier
-  to access lists of those separate parts of the enumeration. Use ``.choices``
-  as a suitable value to pass to :attr:`~Field.choices` in a field definition.
+  ``.choices``, ``.flatchoices``, ``.labels``, ``.values``, and ``.names`` -- to
+  make it easier to access lists of those separate parts of the enumeration.
+  Use ``.choices`` as a suitable value to pass to :attr:`~Field.choices` in a
+  field definition.
 
   .. warning::
 
@@ -297,20 +298,49 @@ you can subclass ``Choices`` and the required concrete data type, e.g.
         APOLLO_16 = 1972, 4, 21, 'Apollo 16 (Orion)'
         APOLLO_17 = 1972, 12, 11, 'Apollo 17 (Challenger)'
 
+.. _field-choices-enum-named-groups:
+
+Enumeration types support :ref:`named groups <field-choices-named-groups>`.
+Use subclasses to denote named groups. A ``__label__`` property is available
+on subclasses if the inferred label is not enough. A ``.flatchoices``
+property is available if you need the flattened version of the choices. The
+``MEDIA_CHOICES`` above can be represented by::
+
+    >>> class Media(models.TextChoices):
+    ...     class Audio(models.TextChoices):
+    ...         VINYL = "vinyl"
+    ...         CD = "cd", "CD"
+    ...     class Video(models.TextChoices):
+    ...         VHS_TAPE = "vhs", "VHS Tape"
+    ...         DVD = "dvd", _("DVD")
+    ...     UNKNOWN = "unknown"
+    >>> Media.choices
+    [('Audio', [('vinyl', 'Vinyl'), ('cd', 'CD')]),
+     ('Video', [('vhs', 'VHS Tape'), ('dvd', 'DVD')]),
+     ('unknown', 'Unknown')]
+    >>> Media.flatchoices
+    [('vinyl', 'Vinyl'), ('cd', 'CD'), ('vhs', 'VHS Tape'),
+     ('dvd', 'DVD'), ('unknown', 'Unknown')]
+
+.. versionchanged:: 4.2
+
+    Support for named groups was added.
+
 There are some additional caveats to be aware of:
 
-- Enumeration types do not support :ref:`named groups
-  <field-choices-named-groups>`.
 - Because an enumeration with a concrete data type requires all values to match
   the type, overriding the :ref:`blank label <field-choices-blank-label>`
   cannot be achieved by creating a member with a value of ``None``. Instead,
-  set the ``__empty__`` attribute on the class::
-
-    class Answer(models.IntegerChoices):
-        NO = 0, _('No')
-        YES = 1, _('Yes')
-
-        __empty__ = _('(Unknown)')
+  set the ``__empty__`` attribute on the class. Note that this member is
+  always the first choice::
+
+    >>> class Answer(models.IntegerChoices):
+    ...     NO = 0, _('No')
+    ...     YES = 1, _('Yes')
+    ... 
+    ...     __empty__ = _('(Unknown)')
+    >>> Answer.choices
+    [(None, '(Unknown)'), (0, 'No'), (1, 'Yes')]
 
 ``db_column``
 -------------

docs/releases/4.2.txt

@@ -184,7 +184,9 @@ Migrations
 Models
 ~~~~~~
 
-* ...
+* :ref:`Enumeration types <field-choices-enum-types>` such as ``TextChoices``,
+  ``IntegerChoices``, and ``Choices`` used to define :attr:`.Field.choices` now
+  support :ref:`named groups <field-choices-enum-named-groups>`.
 
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~

tests/model_enums/tests.py

@@ -41,6 +41,35 @@ class Gender(models.TextChoices):
     __empty__ = "(Undeclared)"
 
 
+class Media(models.TextChoices):
+    class Audio(models.TextChoices):
+        VINYL = "vinyl"
+        CD = "cd", "CD"
+
+    class Video(models.TextChoices):
+        VHS_TAPE = "vhs", "VHS Tape"
+        DVD = "dvd", _("DVD")
+
+    UNKNOWN = "unknown"
+
+
+class Favorite(models.IntegerChoices):
+    NOTHING = 1
+
+    class Colors(models.IntegerChoices):
+        RED = 2
+        GREEN = 3
+        BLUE = 4
+
+    class FoodGroup(models.IntegerChoices):
+        PIZZA = 5
+        ICE_CREAM = 6
+
+        __label__ = "Food"
+
+    __empty__ = _("(Unknown)")
+
+
 class ChoicesTests(SimpleTestCase):
     def test_integerchoices(self):
         self.assertEqual(
@@ -146,6 +175,91 @@ class BlankStr(models.TextChoices):
         self.assertEqual(BlankStr.values, ["", "ONE"])
         self.assertEqual(BlankStr.names, ["EMPTY", "ONE"])
 
+    def test_namedgroups_integerchoices(self):
+        self.assertEqual(
+            Favorite.choices,
+            [
+                (None, "(Unknown)"),
+                (1, "Nothing"),
+                (
+                    "Colors",
+                    [
+                        (2, "Red"),
+                        (3, "Green"),
+                        (4, "Blue"),
+                    ],
+                ),
+                (
+                    "Food",
+                    [
+                        (5, "Pizza"),
+                        (6, "Ice Cream"),
+                    ],
+                ),
+            ],
+        )
+        self.assertEqual(
+            Favorite.labels,
+            ["(Unknown)", "Nothing", "Red", "Green", "Blue", "Pizza", "Ice Cream"],
+        )
+        self.assertEqual(Favorite.values, [None, 1, 2, 3, 4, 5, 6])
+        self.assertEqual(
+            Favorite.names,
+            ["__empty__", "NOTHING", "RED", "GREEN", "BLUE", "PIZZA", "ICE_CREAM"],
+        )
+
+        self.assertEqual(repr(Favorite.RED), "Favorite.RED")
+        self.assertEqual(Favorite.RED.label, "Red")
+        self.assertEqual(Favorite.RED.value, 2)
+        self.assertEqual(Favorite["RED"], Favorite.RED)
+        self.assertEqual(Favorite(2), Favorite.RED)
+
+        self.assertIsInstance(Favorite, type(models.Choices))
+        self.assertIsInstance(Favorite.RED, Favorite)
+        self.assertIsInstance(Favorite.RED.value, int)
+
+    def test_namedgroups_textchoices(self):
+        self.assertEqual(
+            Media.choices,
+            [
+                (
+                    "Audio",
+                    [
+                        ("vinyl", "Vinyl"),
+                        ("cd", "CD"),
+                    ],
+                ),
+                (
+                    "Video",
+                    [
+                        ("vhs", "VHS Tape"),
+                        ("dvd", "DVD"),
+                    ],
+                ),
+                ("unknown", "Unknown"),
+            ],
+        )
+        self.assertEqual(
+            Media.labels,
+            ["Vinyl", "CD", "VHS Tape", "DVD", "Unknown"],
+        )
+        self.assertEqual(Media.values, ["vinyl", "cd", "vhs", "dvd", "unknown"])
+        self.assertEqual(
+            Media.names,
+            ["VINYL", "CD", "VHS_TAPE", "DVD", "UNKNOWN"],
+        )
+
+        self.assertEqual(repr(Media.DVD), "Media.DVD")
+        self.assertEqual(Media.DVD.label, "DVD")
+        self.assertEqual(Media.DVD.value, "dvd")
+        self.assertEqual(Media["DVD"], Media.DVD)
+        self.assertEqual(Media("dvd"), Media.DVD)
+
+        self.assertIsInstance(Media, type(models.Choices))
+        self.assertIsInstance(Media.DVD, Media)
+        self.assertIsInstance(Media.DVD.label, Promise)
+        self.assertIsInstance(Media.DVD.value, str)
+
     def test_invalid_definition(self):
         msg = "'str' object cannot be interpreted as an integer"
         with self.assertRaisesMessage(TypeError, msg):
@@ -161,8 +275,26 @@ class Fruit(models.IntegerChoices):
                 APPLE = 1, "Apple"
                 PINEAPPLE = 1, "Pineapple"
 
+        msg = "Attempted to reuse key: 'ONE'"
+        with self.assertRaisesMessage(TypeError, msg):
+
+            class NamedGroup1(models.IntegerChoices):
+                ONE = 1
+
+                class Group1(models.IntegerChoices):
+                    ONE = 1
+
+        msg = "duplicate values found in <enum 'NamedGroup2'>: B -> A"
+        with self.assertRaisesMessage(ValueError, msg):
+
+            class NamedGroup2(models.IntegerChoices):
+                A = 1
+
+                class Group2(models.IntegerChoices):
+                    B = 1
+
     def test_str(self):
-        for test in [Gender, Suit, YearInSchool, Vehicle]:
+        for test in [Gender, Suit, YearInSchool, Vehicle, Media, Favorite]:
             for member in test:
                 with self.subTest(member=member):
                     self.assertEqual(str(test[member.name]), str(member.value))
@@ -183,6 +315,13 @@ def test_label_member(self):
         self.assertEqual(Stationery.label.value, "label")
         self.assertEqual(Stationery.label.name, "label")
 
+    def test_named_group_member(self):
+        # named_group can be used as a member.
+        Groups = models.TextChoices("Groups", "named_group big small")
+        self.assertEqual(Groups.named_group.label, "Named Group")
+        self.assertEqual(Groups.named_group.value, "named_group")
+        self.assertEqual(Groups.named_group.name, "named_group")
+
     def test_do_not_call_in_templates_member(self):
         # do_not_call_in_templates is not implicitly treated as a member.
         Special = models.IntegerChoices("Special", "do_not_call_in_templates")