Coverage for utilities/choices.py: 90%

77 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-10-10 18:35 +0000

1import enum 

2from typing import Any 

3 

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 _ 

8 

9from utilities.data import get_config_value_ci 

10from utilities.string import enum_key 

11 

12__all__ = ( 

13 'Choice', 

14 'ChoiceSet', 

15 'unpack_grouped_choices', 

16) 

17 

18 

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 

30 

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 

38 

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 

43 

44 def __repr__(self): 

45 return f'Choice(value={self.value!r}, label={self.label!r})' 

46 

47 

48class ChoiceSetMeta(type): 

49 """ 

50 Metaclass for ChoiceSet 

51 """ 

52 def __new__(mcs, name, bases, attrs): 

53 

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) 

70 

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'] = {} 

76 

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 

92 

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)) 

100 

101 return super().__new__(mcs, name, bases, attrs) 

102 

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', ()) 

106 

107 def __iter__(cls): 

108 return iter(getattr(cls, '_choices', ())) 

109 

110 

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() 

119 

120 @classmethod 

121 def values(cls): 

122 return [c[0] for c in unpack_grouped_choices(cls._choices)] 

123 

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) 

134 

135 

136def unpack_grouped_choices(choices): 

137 """ 

138 Unpack a grouped choices hierarchy into a flat list of two-tuples. For example: 

139 

140 choices = ( 

141 ('Foo', ( 

142 (1, 'A'), 

143 (2, 'B') 

144 )), 

145 ('Bar', ( 

146 (3, 'C'), 

147 (4, 'D') 

148 )) 

149 ) 

150 

151 becomes: 

152 

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