From 27ac5a368091823f071eed198f6125eee2de8dfa Mon Sep 17 00:00:00 2001
From: Tom Christie <tom@tomchristie.com>
Date: Thu, 6 Aug 2015 11:43:03 +0100
Subject: [PATCH] Support grouped choices

---
 rest_framework/fields.py                      | 87 ++++++++++++++-----
 .../rest_framework/horizontal/select.html     | 10 ++-
 .../horizontal/select_multiple.html           | 10 ++-
 .../rest_framework/inline/select.html         | 11 ++-
 .../inline/select_multiple.html               | 12 ++-
 .../rest_framework/vertical/select.html       | 11 ++-
 .../vertical/select_multiple.html             | 10 ++-
 rest_framework/utils/field_mapping.py         |  4 +-
 tests/test_model_serializer.py                |  2 +-
 9 files changed, 118 insertions(+), 39 deletions(-)

diff --git a/rest_framework/fields.py b/rest_framework/fields.py
index 9e44a3dc3..4aa33452f 100644
--- a/rest_framework/fields.py
+++ b/rest_framework/fields.py
@@ -5,7 +5,6 @@ import copy
 import datetime
 import decimal
 import inspect
-import itertools
 import re
 import uuid
 
@@ -109,32 +108,54 @@ def set_value(dictionary, keys, value):
     dictionary[keys[-1]] = value
 
 
-def flatten_choice(choice):
+def pairwise_choices(choices):
     """
-    Convert a single choices choice into a flat list of choices.
+    Convert any single choices into key/value pairs instead.
 
-    Returns a list of choices pairs.
-
-    flatten_choice(1) -> [(1, 1)]
-    flatten_choice((1, '1st')) -> [(1, '1st')]
-    flatten_choice(('Grp', ((1, '1st'), (2, '2nd')))) -> [(1, '1st'), (2, '2nd')]
+    pairwise_choices([1]) -> [(1, 1)]
+    pairwise_choices([(1, '1st')]) -> [(1, '1st')]
+    pairwise_choices([('Group', ((1, '1st'), 2))]) -> [(1, '1st'), (2, 2)]
     """
     # Allow single, paired or grouped choices style:
     # choices = [1, 2, 3]
     # choices = [(1, 'First'), (2, 'Second'), (3, 'Third')]
     # choices = [('Category', ((1, 'First'), (2, 'Second'))), (3, 'Third')]
-    if (not isinstance(choice, (list, tuple))):
-        # single choice
-        return [(choice, choice)]
-    else:
-        key, display_value = choice
-        if isinstance(display_value, (list, tuple)):
-            # grouped choices
-            sub_choices = [flatten_choice(c) for c in display_value]
-            return list(itertools.chain(*sub_choices))
+    ret = []
+    for choice in choices:
+        if (not isinstance(choice, (list, tuple))):
+            # single choice
+            item = (choice, choice)
+            ret.append(item)
         else:
-            # paired choice
-            return [(key, display_value)]
+            key, value = choice
+            if isinstance(value, (list, tuple)):
+                # grouped choices (category, sub choices)
+                item = (key, pairwise_choices(value))
+                ret.append(item)
+            else:
+                # paired choice (key, display value)
+                item = (key, value)
+                ret.append(item)
+    return ret
+
+
+def flatten_choices(choices):
+    """
+    Convert a group choices into a flat list of choices.
+
+    flatten_choices([(1, '1st')]) -> [(1, '1st')]
+    flatten_choices([('Grp', ((1, '1st'), (2, '2nd')))]) -> [(1, '1st'), (2, '2nd')]
+    """
+    ret = []
+    for key, value in choices:
+        if isinstance(value, (list, tuple)):
+            # grouped choices (category, sub choices)
+            ret.extend(flatten_choices(value))
+        else:
+            # choice (key, display value)
+            item = (key, value)
+            ret.append(item)
+    return ret
 
 
 class CreateOnlyDefault(object):
@@ -1140,8 +1161,8 @@ class ChoiceField(Field):
     }
 
     def __init__(self, choices, **kwargs):
-        flat_choices = [flatten_choice(c) for c in choices]
-        self.choices = OrderedDict(list(itertools.chain(*flat_choices)))
+        self.grouped_choices = pairwise_choices(choices)
+        self.choices = OrderedDict(flatten_choices(self.grouped_choices))
 
         # Map the string representation of choices to the underlying value.
         # Allows us to deal with eg. integer choices while supporting either
@@ -1168,6 +1189,30 @@ class ChoiceField(Field):
             return value
         return self.choice_strings_to_values.get(six.text_type(value), value)
 
+    def iter_options(self):
+        class StartOptionGroup(object):
+            start_option_group = True
+
+            def __init__(self, label):
+                self.label = label
+
+        class EndOptionGroup(object):
+            end_option_group = True
+
+        class Option(object):
+            def __init__(self, value, display_text):
+                self.value = value
+                self.display_text = display_text
+
+        for key, value in self.grouped_choices:
+            if isinstance(value, (list, tuple)):
+                yield StartOptionGroup(label=key)
+                for sub_key, sub_value in value:
+                    yield Option(value=sub_key, display_text=sub_value)
+                yield EndOptionGroup()
+            else:
+                yield Option(value=key, display_text=value)
+
 
 class MultipleChoiceField(ChoiceField):
     default_error_messages = {
diff --git a/rest_framework/templates/rest_framework/horizontal/select.html b/rest_framework/templates/rest_framework/horizontal/select.html
index 8957eef9f..4aa3db3de 100644
--- a/rest_framework/templates/rest_framework/horizontal/select.html
+++ b/rest_framework/templates/rest_framework/horizontal/select.html
@@ -10,8 +10,14 @@
       {% if field.allow_null or field.allow_blank %}
         <option value="" {% if not field.value %}selected{% endif %}>--------</option>
       {% endif %}
-      {% for key, text in field.choices.items %}
-        <option value="{{ key }}" {% if key == field.value %}selected{% endif %}>{{ text }}</option>
+      {% for select in field.iter_options %}
+          {% if select.start_option_group %}
+            <optgroup label="{{ select.label }}">
+          {% elif select.end_option_group %}
+            </optgroup>
+          {% else %}
+            <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+          {% endif %}
       {% endfor %}
     </select>
 
diff --git a/rest_framework/templates/rest_framework/horizontal/select_multiple.html b/rest_framework/templates/rest_framework/horizontal/select_multiple.html
index 407b356ff..b1ca2f143 100644
--- a/rest_framework/templates/rest_framework/horizontal/select_multiple.html
+++ b/rest_framework/templates/rest_framework/horizontal/select_multiple.html
@@ -10,8 +10,14 @@
 
   <div class="col-sm-10">
     <select multiple {{ field.choices|yesno:",disabled" }} class="form-control" name="{{ field.name }}">
-      {% for key, text in field.choices.items %}
-        <option value="{{ key }}" {% if key in field.value %}selected{% endif %}>{{ text }}</option>
+      {% for select in field.iter_options %}
+        {% if select.start_option_group %}
+          <optgroup label="{{ select.label }}">
+        {% elif select.end_option_group %}
+          </optgroup>
+        {% else %}
+          <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+        {% endif %}
       {% empty %}
           <option>{{ no_items }}</option>
       {% endfor %}
diff --git a/rest_framework/templates/rest_framework/inline/select.html b/rest_framework/templates/rest_framework/inline/select.html
index c3b2923a1..d5204bb82 100644
--- a/rest_framework/templates/rest_framework/inline/select.html
+++ b/rest_framework/templates/rest_framework/inline/select.html
@@ -9,9 +9,14 @@
     {% if field.allow_null or field.allow_blank %}
       <option value="" {% if not field.value %}selected{% endif %}>--------</option>
     {% endif %}
-
-    {% for key, text in field.choices.items %}
-      <option value="{{ key }}" {% if key == field.value %}selected{% endif %}>{{ text }}</option>
+    {% for select in field.iter_options %}
+        {% if select.start_option_group %}
+          <optgroup label="{{ select.label }}">
+        {% elif select.end_option_group %}
+          </optgroup>
+        {% else %}
+          <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+        {% endif %}
     {% endfor %}
   </select>
 </div>
diff --git a/rest_framework/templates/rest_framework/inline/select_multiple.html b/rest_framework/templates/rest_framework/inline/select_multiple.html
index b9b8bb4ea..3bae43eb8 100644
--- a/rest_framework/templates/rest_framework/inline/select_multiple.html
+++ b/rest_framework/templates/rest_framework/inline/select_multiple.html
@@ -9,9 +9,15 @@
   {% endif %}
 
   <select multiple {{ field.choices|yesno:",disabled" }} class="form-control" name="{{ field.name }}">
-    {% for key, text in field.choices.items %}
-      <option value="{{ key }}" {% if key in field.value %}selected{% endif %}>{{ text }}</option>
-    {% empty %}
+      {% for select in field.iter_options %}
+          {% if select.start_option_group %}
+            <optgroup label="{{ select.label }}">
+          {% elif select.end_option_group %}
+            </optgroup>
+          {% else %}
+            <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+          {% endif %}
+      {% empty %}
       <option>{{ no_items }}</option>
     {% endfor %}
   </select>
diff --git a/rest_framework/templates/rest_framework/vertical/select.html b/rest_framework/templates/rest_framework/vertical/select.html
index dc32d39dd..ce30022d8 100644
--- a/rest_framework/templates/rest_framework/vertical/select.html
+++ b/rest_framework/templates/rest_framework/vertical/select.html
@@ -9,9 +9,14 @@
     {% if field.allow_null or field.allow_blank %}
       <option value="" {% if not field.value %}selected{% endif %}>--------</option>
     {% endif %}
-
-    {% for key, text in field.choices.items %}
-      <option value="{{ key }}" {% if key == field.value %}selected{% endif %}>{{ text }}</option>
+    {% for select in field.iter_options %}
+        {% if select.start_option_group %}
+          <optgroup label="{{ select.label }}">
+        {% elif select.end_option_group %}
+          </optgroup>
+        {% else %}
+          <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+        {% endif %}
     {% endfor %}
   </select>
 
diff --git a/rest_framework/templates/rest_framework/vertical/select_multiple.html b/rest_framework/templates/rest_framework/vertical/select_multiple.html
index 2bb3d5ae4..cb65cec52 100644
--- a/rest_framework/templates/rest_framework/vertical/select_multiple.html
+++ b/rest_framework/templates/rest_framework/vertical/select_multiple.html
@@ -9,8 +9,14 @@
   {% endif %}
 
   <select multiple {{ field.choices|yesno:",disabled" }} class="form-control" name="{{ field.name }}">
-    {% for key, text in field.choices.items %}
-      <option value="{{ key }}" {% if key in field.value %}selected{% endif %}>{{ text }}</option>
+    {% for select in field.iter_options %}
+        {% if select.start_option_group %}
+          <optgroup label="{{ select.label }}">
+        {% elif select.end_option_group %}
+          </optgroup>
+        {% else %}
+          <option value="{{ select.value }}" {% if select.value == field.value %}selected{% endif %}>{{ select.display_text }}</option>
+        {% endif %}
     {% empty %}
         <option>{{ no_items }}</option>
     {% endfor %}
diff --git a/rest_framework/utils/field_mapping.py b/rest_framework/utils/field_mapping.py
index 917c8458e..cebbcc252 100644
--- a/rest_framework/utils/field_mapping.py
+++ b/rest_framework/utils/field_mapping.py
@@ -107,10 +107,10 @@ def get_field_kwargs(field_name, model_field):
                               isinstance(model_field, models.TextField)):
         kwargs['allow_blank'] = True
 
-    if model_field.flatchoices:
+    if model_field.choices:
         # If this model field contains choices, then return early.
         # Further keyword arguments are not valid.
-        kwargs['choices'] = model_field.flatchoices
+        kwargs['choices'] = model_field.choices
         return kwargs
 
     # Ensure that max_length is passed explicitly as a keyword arg,
diff --git a/tests/test_model_serializer.py b/tests/test_model_serializer.py
index 8f3f82b63..00c6319ea 100644
--- a/tests/test_model_serializer.py
+++ b/tests/test_model_serializer.py
@@ -181,7 +181,7 @@ class TestRegularFieldMappings(TestCase):
                 null_field = IntegerField(allow_null=True, required=False)
                 default_field = IntegerField(required=False)
                 descriptive_field = IntegerField(help_text='Some help text', label='A label')
-                choices_field = ChoiceField(choices=[('red', 'Red'), ('blue', 'Blue'), ('green', 'Green')])
+                choices_field = ChoiceField(choices=(('red', 'Red'), ('blue', 'Blue'), ('green', 'Green')))
         """)
         if six.PY2:
             # This particular case is too awkward to resolve fully across