From 342034b072b584b7f0c738f455d71b8979a8d85f Mon Sep 17 00:00:00 2001
From: Parbhat Puri <parbhatpuri17@gmail.com>
Date: Mon, 15 Jul 2019 14:46:29 +0000
Subject: [PATCH] Group meta categories choices by parent

---
 opentech/apply/funds/forms.py   | 41 ++++++++++++++++++++++++++++-----
 opentech/apply/funds/widgets.py |  2 +-
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/opentech/apply/funds/forms.py b/opentech/apply/funds/forms.py
index ba176028a..d04faa70e 100644
--- a/opentech/apply/funds/forms.py
+++ b/opentech/apply/funds/forms.py
@@ -1,3 +1,9 @@
+from functools import partial
+from itertools import groupby
+from operator import attrgetter, methodcaller
+
+from django.forms.models import ModelChoiceIterator
+
 from django import forms
 from django.utils.text import mark_safe, slugify
 from django.utils.translation import ugettext_lazy as _
@@ -289,18 +295,41 @@ class UpdatePartnersForm(forms.ModelForm):
         return instance
 
 
-class MetaCategoryMultipleChoiceField(forms.ModelMultipleChoiceField):
+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 GroupedModelMultipleChoiceField(forms.ModelMultipleChoiceField):
+    def __init__(self, *args, choices_groupby, **kwargs):
+        if isinstance(choices_groupby, str):
+            choices_groupby = methodcaller(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)
+
     def label_from_instance(self, obj):
-        depth_line = '-' * (obj.get_depth() - 2)
-        label = "{} {}".format(depth_line, super().label_from_instance(obj))
-        return {'label': label, 'disabled': not obj.is_leaf()}
+        return {'label': super().label_from_instance(obj), 'disabled': not obj.is_leaf()}
 
 
 class UpdateMetaCategoriesForm(forms.ModelForm):
-    meta_categories = MetaCategoryMultipleChoiceField(
+    meta_categories = GroupedModelMultipleChoiceField(
         queryset=None,  # updated in init method
         widget=MetaCategorySelect2Widget(attrs={'data-placeholder': 'Meta categories'}),
         label='Meta categories',
+        choices_groupby='get_parent',
         required=False,
         help_text='Meta categories are hierarchical in nature, - highlights the depth in the tree.',
     )
@@ -312,4 +341,4 @@ class UpdateMetaCategoriesForm(forms.ModelForm):
     def __init__(self, *args, **kwargs):
         kwargs.pop('user')
         super().__init__(*args, **kwargs)
-        self.fields['meta_categories'].queryset = MetaCategory.get_root_descendants()
+        self.fields['meta_categories'].queryset = MetaCategory.get_root_descendants().exclude(depth=2)
diff --git a/opentech/apply/funds/widgets.py b/opentech/apply/funds/widgets.py
index d17e7bcbb..4af3d76a7 100644
--- a/opentech/apply/funds/widgets.py
+++ b/opentech/apply/funds/widgets.py
@@ -22,7 +22,7 @@ class Select2MultiCheckboxesWidget(Select2MultipleWidget):
         return attrs
 
 
-class MetaCategorySelect2Widget(Select2MultiCheckboxesWidget):
+class MetaCategorySelect2Widget(Select2MultipleWidget):
 
     def create_option(self, name, value, label, selected, index, subindex=None, attrs=None):
         disabled = False
-- 
GitLab