1 # -*- coding: utf-8 -*-
 
   2 from django import forms
 
   3 from django.forms.models import ModelChoiceIterator
 
   4 from django.utils.translation import ugettext as _
 
   7 from catalogue.models import Lesson
 
  10 class GroupedModelChoiceIterator(ModelChoiceIterator):
 
  11     def __init__(self, field):
 
  12         super(GroupedModelChoiceIterator, self).__init__(field)
 
  13         self.queryset = self.field.grouping_model.objects
 
  14         self.items_queryset = self.field.queryset
 
  16     def choice(self, obj):
 
  17         items_query = self.items_queryset.filter(**{self.field.grouping_fk_field: obj})
 
  18         items = [super(GroupedModelChoiceIterator, self).choice(item) for item in items_query.all()]
 
  19         return unicode(obj), items
 
  22 class GroupedModelChoiceField(forms.ModelChoiceField):
 
  24     def __init__(self, queryset, grouping_fk_field, **kwargs):
 
  25         self.grouping_fk_field = grouping_fk_field
 
  26         self.grouping_model = queryset.model._meta.get_field(grouping_fk_field).rel.to
 
  27         super(GroupedModelChoiceField, self).__init__(queryset, **kwargs)
 
  29     def _get_choices(self):
 
  30         toret = super(GroupedModelChoiceField, self)._get_choices()
 
  31         if isinstance(toret, ModelChoiceIterator):
 
  32             toret = GroupedModelChoiceIterator(self)
 
  35     choices = property(_get_choices, forms.ModelChoiceField.choices.fset)
 
  38 class PostForm(pybb.forms.PostForm):
 
  39     lesson = GroupedModelChoiceField(
 
  40         label=_('Related lesson'), queryset=Lesson.objects.all(),
 
  41         grouping_fk_field='section', required=False)