Support nested choice groups

netbox/utilities/choices.py

@@ -15,21 +15,21 @@ class ChoiceSetMeta(type):
             except KeyError:
                 pass
 
-        # Define choice tuples
+        # Define choice tuples and color maps
-        # TODO: Support optgroup nesting
+        attrs['_choices'] = []
-        attrs['_choices'] = [
+        attrs['colors'] = {}
-            (c[0], c[1]) for c in attrs['CHOICES']
+        for choice in attrs['CHOICES']:
-        ]
+            if isinstance(choice[1], (list, tuple)):
-
+                grouped_choices = []
-        # Define color maps
+                for c in choice[1]:
-        # TODO: Support optgroup nesting
+                    grouped_choices.append((c[0], c[1]))
-        colors = {}
+                    if len(c) == 3:
-        for c in attrs['CHOICES']:
+                        attrs['colors'][c[0]] = c[2]
-            try:
+                attrs['_choices'].append((choice[0], grouped_choices))
-                colors[c[0]] = c[2]
+            else:
-            except IndexError:
+                attrs['_choices'].append((choice[0], choice[1]))
-                pass
+                if len(choice) == 3:
-        attrs['colors'] = colors
+                    attrs['colors'][choice[0]] = choice[2]
 
         return super().__new__(mcs, name, bases, attrs)
 
@@ -48,12 +48,12 @@ class ChoiceSet(metaclass=ChoiceSetMeta):
 
     @classmethod
     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_dict(cls):
         # Unpack grouped choices before casting as a dict
-        return dict(unpack_grouped_choices(cls.CHOICES))
+        return dict(unpack_grouped_choices(cls._choices))
 
 
 def unpack_grouped_choices(choices):