diff --git a/app/apps/common/fields/forms/grouped_select.py b/app/apps/common/fields/forms/grouped_select.py new file mode 100644 index 0000000..6e9c4e2 --- /dev/null +++ b/app/apps/common/fields/forms/grouped_select.py @@ -0,0 +1,49 @@ +from django import forms +from django.forms.models import ModelChoiceIterator +from django.utils.translation import gettext_lazy as _ + + +class GroupedModelChoiceIterator(ModelChoiceIterator): + def __init__(self, field, group_by): + self.group_by = group_by + super().__init__(field) + + def __iter__(self): + if self.field.empty_label is not None: + yield "", self.field.empty_label + + queryset = self.queryset + grouped_queryset = {} + + for obj in queryset: + group_key = getattr(obj, self.group_by, None) if self.group_by else None + grouped_queryset.setdefault(group_key, []).append(obj) + + for group, items in grouped_queryset.items(): + yield ( + group if group is not None else _("Ungrouped"), + [ + (self.field.prepare_value(obj), self.field.label_from_instance(obj)) + for obj in items + ], + ) + + +class GroupedModelChoiceField(forms.ModelChoiceField): + def __init__(self, *args, group_by=None, **kwargs): + self.group_by = group_by + super().__init__(*args, **kwargs) + + @property + def choices(self): + return GroupedModelChoiceIterator(self, self.group_by) + + +class GroupedModelMultipleChoiceField(forms.ModelMultipleChoiceField): + def __init__(self, *args, group_by=None, **kwargs): + self.group_by = group_by + super().__init__(*args, **kwargs) + + @property + def choices(self): + return GroupedModelChoiceIterator(self, self.group_by)