from django.conf import settings from django.utils.translation import gettext_lazy as _ __all__ = ( 'ChoiceSet', 'unpack_grouped_choices', ) class ChoiceSetMeta(type): """ Metaclass for ChoiceSet """ def __new__(mcs, name, bases, attrs): # Extend static choices with any configured choices if key := attrs.get('key'): assert type(attrs['CHOICES']) is list, _( "{name} has a key defined but CHOICES is not a list" ).format(name=name) app = attrs['__module__'].split('.', 1)[0] replace_key = f'{app}.{key}' extend_key = f'{replace_key}+' if replace_key else None if replace_key and replace_key in settings.FIELD_CHOICES: # Replace the stock choices attrs['CHOICES'] = settings.FIELD_CHOICES[replace_key] elif extend_key and extend_key in settings.FIELD_CHOICES: # Extend the stock choices attrs['CHOICES'].extend(settings.FIELD_CHOICES[extend_key]) # Define choice tuples and color maps attrs['_choices'] = [] attrs['colors'] = {} for choice in attrs['CHOICES']: if isinstance(choice[1], (list, tuple)): grouped_choices = [] for c in choice[1]: grouped_choices.append((c[0], c[1])) if len(c) == 3: attrs['colors'][c[0]] = c[2] attrs['_choices'].append((choice[0], grouped_choices)) else: attrs['_choices'].append((choice[0], choice[1])) if len(choice) == 3: attrs['colors'][choice[0]] = choice[2] return super().__new__(mcs, name, bases, attrs) def __call__(cls, *args, **kwargs): # django-filters will check if a 'choices' value is callable, and if so assume that it returns an iterable return getattr(cls, '_choices', ()) def __iter__(cls): return iter(getattr(cls, '_choices', ())) class ChoiceSet(metaclass=ChoiceSetMeta): """ Holds an iterable of choice tuples suitable for passing to a Django model or form field. Choices can be defined statically within the class as CHOICES and/or gleaned from the FIELD_CHOICES configuration parameter. """ CHOICES = list() @classmethod def values(cls): return [c[0] for c in unpack_grouped_choices(cls._choices)] def unpack_grouped_choices(choices): """ Unpack a grouped choices hierarchy into a flat list of two-tuples. For example: choices = ( ('Foo', ( (1, 'A'), (2, 'B') )), ('Bar', ( (3, 'C'), (4, 'D') )) ) becomes: choices = ( (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D') ) """ unpacked_choices = [] for key, value in choices: if isinstance(value, (list, tuple)): # Entered an optgroup for optgroup_key, optgroup_value in value: unpacked_choices.append((optgroup_key, optgroup_value)) else: unpacked_choices.append((key, value)) return unpacked_choices