Coverage for utilities/choices.py: 90%
77 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 18:35 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 18:35 +0000
1import enum
2from typing import Any
4from django.conf import settings
5from django.core.exceptions import ImproperlyConfigured
6from django.utils.functional import Promise
7from django.utils.translation import gettext_lazy as _
9from utilities.data import get_config_value_ci
10from utilities.string import enum_key
12__all__ = (
13 'Choice',
14 'ChoiceSet',
15 'unpack_grouped_choices',
16)
19class Choice(tuple):
20 """
21 A single choice within a ChoiceSet. Carries the choice's label, and optionally a color and a description, in a
22 single object. A Choice **is** a `(value, label)` two-tuple for backward compatibility with the plain-tuple choice
23 format (so it satisfies `isinstance(choice, tuple)` checks in Django, DRF, etc.); its `color` and `description`
24 are exposed only as attributes.
25 """
26 value: Any
27 label: str | Promise
28 color: str | None
29 description: str | Promise | None
31 def __new__(cls, value, label, color=None, description=None):
32 instance = super().__new__(cls, (value, label))
33 instance.value = value
34 instance.label = label
35 instance.color = color
36 instance.description = description
37 return instance
39 def __getnewargs__(self):
40 # tuple's default reconstructor would call Choice((value, label)) with a single argument, which fails
41 # __new__'s (value, label, ...) signature. Supply the real constructor args so copy/deepcopy/pickle work.
42 return self.value, self.label, self.color, self.description
44 def __repr__(self):
45 return f'Choice(value={self.value!r}, label={self.label!r})'
48class ChoiceSetMeta(type):
49 """
50 Metaclass for ChoiceSet
51 """
52 def __new__(mcs, name, bases, attrs):
54 # Extend static choices with any configured choices
55 if key := attrs.get('key'):
56 if type(attrs['CHOICES']) is not list: 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true
57 raise ImproperlyConfigured(
58 _("{name} has a key defined but CHOICES is not a list").format(name=name)
59 )
60 app = attrs['__module__'].split('.', 1)[0]
61 replace_key = f'{app}.{key}'
62 replace_choices = get_config_value_ci(settings.FIELD_CHOICES, replace_key)
63 if replace_choices is not None: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true
64 attrs['CHOICES'] = replace_choices
65 else:
66 extend_key = f'{replace_key}+'
67 extend_choices = get_config_value_ci(settings.FIELD_CHOICES, extend_key)
68 if extend_choices is not None: 68 ↛ 69line 68 didn't jump to line 69 because the condition on line 68 was never true
69 attrs['CHOICES'].extend(extend_choices)
71 # Build the normalized choice list and the derived color map. Each choice may be defined as a Choice object
72 # (which is preserved as-is so consumers can reference its color/description) or as a plain (value, label) or
73 # (value, label, color) tuple. The colors map is kept for model-level consumers (e.g. get_FOO_color()).
74 attrs['_choices'] = []
75 attrs['colors'] = {}
77 def register(entry):
78 # A choice may be given as a dict (e.g. from FIELD_CHOICES config, to avoid importing Choice), a Choice
79 # object, or a legacy (value, label[, color]) tuple. Dicts and Choices are preserved as Choice objects so
80 # consumers can reference their color/description; legacy tuples are reduced to (value, label). Any color
81 # is also recorded on the colors map for model-level consumers.
82 if isinstance(entry, dict): 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true
83 entry = Choice(**entry)
84 if isinstance(entry, Choice):
85 if entry.color is not None:
86 attrs['colors'][entry.value] = entry.color
87 return entry
88 value, label = entry[0], entry[1]
89 if len(entry) >= 3:
90 attrs['colors'][value] = entry[2]
91 return value, label
93 for choice in attrs['CHOICES']:
94 # A grouped choice is a (group_label, [members]) tuple; Choice and dict entries are always flat choices
95 if not isinstance(choice, (Choice, dict)) and isinstance(choice[1], (list, tuple)):
96 grouped_choices = [register(c) for c in choice[1]]
97 attrs['_choices'].append((choice[0], grouped_choices))
98 else:
99 attrs['_choices'].append(register(choice))
101 return super().__new__(mcs, name, bases, attrs)
103 def __call__(cls, *args, **kwargs):
104 # django-filters will check if a 'choices' value is callable, and if so assume that it returns an iterable
105 return getattr(cls, '_choices', ())
107 def __iter__(cls):
108 return iter(getattr(cls, '_choices', ()))
111class ChoiceSet(metaclass=ChoiceSetMeta):
112 """
113 Holds an iterable of choices suitable for passing to a Django model or form field. Choices can be defined
114 statically within the class as CHOICES and/or gleaned from the FIELD_CHOICES configuration parameter. Each
115 choice may be defined as a Choice object (to carry a color and/or description) or as a plain (value, label) or
116 (value, label, color) tuple.
117 """
118 CHOICES = list()
120 @classmethod
121 def values(cls):
122 return [c[0] for c in unpack_grouped_choices(cls._choices)]
124 @classmethod
125 def as_enum(cls, name=None, prefix=''):
126 """
127 Return the ChoiceSet as an Enum. If no name is provided, "Choices" will be stripped from the class name (if
128 present) and "Enum" will be appended. For example, "CircuitStatusChoices" will become "CircuitStatusEnum".
129 """
130 name = name or f"{cls.__name__.split('Choices')[0]}Enum"
131 prefix = f'{prefix}_' if prefix else ''
132 data = {f'{prefix}{enum_key(v)}'.upper(): v for v in cls.values()}
133 return enum.Enum(name, data)
136def unpack_grouped_choices(choices):
137 """
138 Unpack a grouped choices hierarchy into a flat list of two-tuples. For example:
140 choices = (
141 ('Foo', (
142 (1, 'A'),
143 (2, 'B')
144 )),
145 ('Bar', (
146 (3, 'C'),
147 (4, 'D')
148 ))
149 )
151 becomes:
153 choices = (
154 (1, 'A'),
155 (2, 'B'),
156 (3, 'C'),
157 (4, 'D')
158 )
159 """
160 unpacked_choices = []
161 for key, value in choices:
162 if isinstance(value, (list, tuple)):
163 # Entered an optgroup
164 for optgroup_key, optgroup_value in value:
165 unpacked_choices.append((optgroup_key, optgroup_value))
166 else:
167 unpacked_choices.append((key, value))
168 return unpacked_choices