From 25a95dafd5fa4f784fbe5f3ca5a9d6685f13a31a Mon Sep 17 00:00:00 2001
From: Sander Roosendaal <roosendaalsander@gmail.com>
Date: Wed, 11 Nov 2020 06:43:37 +0100
Subject: [PATCH] courses grouped by country in form

---
 rowers/formfields.py | 31 +++++++++++++++++++++++++++++++
 rowers/models.py     |  7 +++++--
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 rowers/formfields.py

diff --git a/rowers/formfields.py b/rowers/formfields.py
new file mode 100644
index 00000000..3b418ef3
--- /dev/null
+++ b/rowers/formfields.py
@@ -0,0 +1,31 @@
+from functools import partial
+from itertools import groupby
+from operator import attrgetter
+
+from django.forms.models import ModelChoiceIterator, ModelChoiceField
+
+
+class GroupedModelChoiceIterator(ModelChoiceIterator):
+    def __init__(self, field, groupby):
+        self.groupby = groupby
+        super().__init__(field)
+
+    def __iter__(self):
+        if self.field.empty_label is not None:
+            yield ("", self.field.empty_label)
+        queryset = self.queryset
+        # Can't use iterator() when queryset uses prefetch_related()
+        if not queryset._prefetch_related_lookups:
+            queryset = queryset.iterator()
+        for group, objs in groupby(queryset, self.groupby):
+            yield (group, [self.choice(obj) for obj in objs])
+
+
+class GroupedModelChoiceField(ModelChoiceField):
+    def __init__(self, *args, choices_groupby, **kwargs):
+        if isinstance(choices_groupby, str):
+            choices_groupby = attrgetter(choices_groupby)
+        elif not callable(choices_groupby):
+            raise TypeError('choices_groupby must either be a str or a callable accepting a single argument')
+        self.iterator = partial(GroupedModelChoiceIterator, groupby=choices_groupby)
+        super().__init__(*args, **kwargs)
diff --git a/rowers/models.py b/rowers/models.py
index 1dcfd8ce..2b1a13d1 100644
--- a/rowers/models.py
+++ b/rowers/models.py
@@ -69,6 +69,7 @@ tweetapi = twitter.Api(consumer_key=TWEET_CONSUMER_KEY,
                        access_token_secret=TWEET_ACCESS_TOKEN_SECRET)
 
 from rowers.database import *
+from rowers.formfields import *
 
 timezones = (
     (x,x) for x in pytz.common_timezones
@@ -2689,11 +2690,13 @@ class IndoorVirtualRaceForm(ModelForm):
 
 
 class VirtualRaceForm(ModelForm):
-    course = forms.ModelChoiceField(queryset = GeoCourse.objects, empty_label=None)
+    course = GroupedModelChoiceField(
+        queryset = GeoCourse.objects, empty_label=None,
+        choices_groupby='country'
+        )
     registration_closure = forms.SplitDateTimeField(widget=AdminSplitDateTime(),required=False)
     evaluation_closure = forms.SplitDateTimeField(widget=AdminSplitDateTime(),required=True)
 
-
     class Meta:
         model = VirtualRace
         fields = [