A Grouped ModelChoice Field for Django

fields.py

from itertools import groupby
from typing import TYPE_CHECKING, TypeVar, Generic, cast

from django.forms.fields import ChoiceField

from django.forms.models import (
  ModelChoiceIterator, ModelChoiceField, ModelMultipleChoiceField,
)


if TYPE_CHECKING:
  from typing import Any, Callable, Iterator
  from django.db.models import Model
  from django.db.models.query import QuerySet
  from django.utils.choices import (
    CallableChoiceIterator, _ChoicesCallable, _ChoicesInput,
  )

  type GroupedField[_M: Model] = (
    GroupedModelChoiceField[_M] | GroupedModelMultipleChoiceField[_M]
  )
  type GroupedChoicesSetter = (
    _ChoicesInput | _ChoicesCallable | CallableChoiceIterator
  )
  type GroupedChoicesGetter[_M: Model] = (
    _ChoicesInput | CallableChoiceIterator | GroupedModelChoiceIterator[_M]
  )


_M = TypeVar('_M', bound='Model')


class GroupedModelChoiceMixin(Generic[_M]):
  """ModelChoice form field with option groups."""

  group_by_field: str
  group_label: Callable[[object], object]
  _choices: GroupedChoicesSetter

  def __init__(
    self,
    group_by_field: str,
    group_label: Callable[[object], object] | None = None,
    *args: Any,
    **kwargs: Any,
  ) -> None:
    """
    group_by_field is the name of a field on the model
    group_label is a function to return a label for each choice group
    """

    super().__init__(*args, **kwargs)

    self.group_by_field = group_by_field

    if group_label is None:
      def identity(group: object) -> object:
        return group
      self.group_label = identity
    else:
      self.group_label = group_label

  def _get_choices(self) -> GroupedChoicesGetter[_M]:
    """
    Exactly as per ModelChoiceField except returns new iterator class
    """
    if hasattr(self, '_choices'):
      return self._choices
    return GroupedModelChoiceIterator(cast('GroupedField[_M]', self))

  @property
  def choices(self) -> GroupedChoicesGetter[_M]:
    return self._get_choices()

  @choices.setter
  def choices(self, value: GroupedChoicesSetter) -> None:
    getattr(ChoiceField.choices, 'fset')(self, value)


class GroupedModelChoiceField(
  GroupedModelChoiceMixin[_M],
  ModelChoiceField[_M],
):
  pass


class GroupedModelMultipleChoiceField(
  GroupedModelChoiceMixin[_M],
  ModelMultipleChoiceField[_M],
):
  pass


class GroupedModelChoiceIterator(ModelChoiceIterator, Generic[_M]):
  """Iterator for ModelChoices that uses option groups."""

  field: GroupedField[_M]
  queryset: QuerySet[_M]

  def __iter__(self) -> Iterator[Any]:
    if self.field.empty_label is not None:
      yield (u"", self.field.empty_label)

    for group, choices in groupby(
      self.queryset.all(),
      key=lambda row: getattr(row, self.field.group_by_field)
    ):
      yield (self.field.group_label(group), [self.choice(c) for c in choices])