Repository: wtforms/wtforms
Pull request #667 (2020-10-11 23:49:31)

Add support for <optgroup> in SelectField
Add support for <optgroup> in SelectField as described in issue #656 

<!--
Commit checklist:

* add tests that fail without the patch
* ensure all tests pass with ``pytest``
* add documentation to the relevant docstrings or pages
* add ``versionadded`` or ``versionchanged`` directives to relevant docstrings
* add a changelog entry if this patch changes code
-->


davidism: Please rebase this onto master, 2.3.x will not get another feature release.

Patch:
diff --git a/docs/fields.rst b/docs/fields.rst
index fc0c351..3e0e26a 100644
--- a/docs/fields.rst
+++ b/docs/fields.rst
@@ -274,13 +274,16 @@ refer to a single input from the form.
 
 .. class:: SelectField(default field arguments, choices=[], coerce=str, option_widget=None, validate_choice=True)
 
-    Select fields take a ``choices`` parameter which is either a list of
-    ``(value, label)`` pairs, or a function taking no argument, and returning
-    a list of ``(value, label)`` pairs. It can also be a list of only values, in
-    which case the value is used as the label. The value can be any
-    type, but because form data is sent to the browser as strings, you
-    will need to provide a ``coerce`` function that converts a string
-    back to the expected type.
+    Select fields take a ``choices`` parameter which is either:
+
+    * a list of ``(value, label)`` pairs. It can also be a list of only values, in
+      which case the value is used as the label. The value can be of any
+      type, but because form data is sent to the browser as strings, you
+      will need to provide a ``coerce`` function that converts a string
+      back to the expected type.
+    * a dictionary of ``{label: list}`` pairs defining groupings of options.
+    * a function taking no argument, and returning either a list or a dictionary.
+
 
     **Select fields with static choice values**::
 
diff --git a/src/wtforms/fields/core.py b/src/wtforms/fields/core.py
index 96484a9..2b04972 100644
--- a/src/wtforms/fields/core.py
+++ b/src/wtforms/fields/core.py
@@ -498,6 +498,12 @@ class SelectFieldBase(Field):
         """
         raise NotImplementedError()
 
+    def has_groups(self):
+        return False
+
+    def iter_groups(self):
+        raise NotImplementedError()
+
     def __iter__(self):
         opts = dict(
             widget=self.option_widget,
@@ -536,18 +542,40 @@ class SelectField(SelectFieldBase):
         self.coerce = coerce
         if callable(choices):
             choices = choices()
-        self.choices = list(choices) if choices is not None else None
+        if choices is not None:
+            self.choices = choices if isinstance(choices, dict) else list(choices)
+        else:
+            self.choices = None
         self.validate_choice = validate_choice
 
     def iter_choices(self):
         if not self.choices:
             choices = []
-        elif isinstance(self.choices[0], (list, tuple)):
+        elif isinstance(self.choices, dict):
+            choices = list(itertools.chain.from_iterable(self.choices.values()))
+        else:
             choices = self.choices
+
+        return self._choices_generator(choices)
+
+    def has_groups(self):
+        return isinstance(self.choices, dict)
+
+    def iter_groups(self):
+        if isinstance(self.choices, dict):
+            for label, choices in self.choices.items():
+                yield (label, self._choices_generator(choices))
+
+    def _choices_generator(self, choices):
+        if choices:
+            if isinstance(choices[0], (list, tuple)):
+                _choices = choices
+            else:
+                _choices = zip(choices, choices)
         else:
-            choices = zip(self.choices, self.choices)
+            _choices = []
 
-        for value, label in choices:
+        for value, label in _choices:
             yield (value, label, self.coerce(value) == self.data)
 
     def process_data(self, value):
@@ -589,15 +617,16 @@ class SelectMultipleField(SelectField):
 
     widget = widgets.Select(multiple=True)
 
-    def iter_choices(self):
-        if not self.choices:
-            choices = []
-        elif isinstance(self.choices[0], (list, tuple)):
-            choices = self.choices
+    def _choices_generator(self, choices):
+        if choices:
+            if isinstance(choices[0], (list, tuple)):
+                _choices = choices
+            else:
+                _choices = zip(choices, choices)
         else:
-            choices = zip(self.choices, self.choices)
+            _choices = []
 
-        for value, label in choices:
+        for value, label in _choices:
             selected = self.data is not None and self.coerce(value) in self.data
             yield (value, label, selected)
 
diff --git a/src/wtforms/widgets/core.py b/src/wtforms/widgets/core.py
index a0e0e61..f5ca315 100644
--- a/src/wtforms/widgets/core.py
+++ b/src/wtforms/widgets/core.py
@@ -321,6 +321,10 @@ class Select:
     The field must provide an `iter_choices()` method which the widget will
     call on rendering; this method must yield tuples of
     `(value, label, selected)`.
+    It also must provide a `has_groups()` method which tells whether choices
+    are divided into groups, and if they do, the field must have an
+    `iter_groups()` method that yields tuples of `(label, choices)`, where
+    `choices` is a iterable of `(value, label, selected)` tuples.
     """
 
     validation_attrs = ["required"]
@@ -337,8 +341,15 @@ class Select:
             if k in self.validation_attrs and k not in kwargs:
                 kwargs[k] = getattr(flags, k)
         html = ["<select %s>" % html_params(name=field.name, **kwargs)]
-        for val, label, selected in field.iter_choices():
-            html.append(self.render_option(val, label, selected))
+        if field.has_groups():
+            for group, choices in field.iter_groups():
+                html.append("<optgroup %s>" % html_params(label=group))
+                for val, label, selected in choices:
+                    html.append(self.render_option(val, label, selected))
+                html.append("</optgroup>")
+        else:
+            for val, label, selected in field.iter_choices():
+                html.append(self.render_option(val, label, selected))
         html.append("</select>")
         return Markup("".join(html))
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index d7dc913..3793398 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -110,6 +110,12 @@ class DummyField:
     def iter_choices(self):
         return iter(self.data)
 
+    def iter_groups(self):
+        return []
+
+    def has_groups(self):
+        return False
+
     def gettext(self, string):
         return self._translations.gettext(string)
 
diff --git a/tests/fields/test_select.py b/tests/fields/test_select.py
index c411f4a..dd0d263 100644
--- a/tests/fields/test_select.py
+++ b/tests/fields/test_select.py
@@ -146,7 +146,7 @@ def test_choice_shortcut_post():
     assert len(form.a.errors) == 0
 
 
-@pytest.mark.parametrize("choices", [[], None])
+@pytest.mark.parametrize("choices", [[], None, {}])
 def test_empty_choice(choices):
     F = make_form(a=SelectField(choices=choices, validate_choice=False))
     form = F(a="bar")
@@ -208,3 +208,36 @@ def test_render_kw_preserved():
         '<option value="bar">bar</option>'
         "</select>"
     )
+
+
+def test_optgroup():
+    F = make_form(a=SelectField(choices={"hello": [("a", "Foo")]}))
+    form = F(a="a")
+
+    assert (
+        '<optgroup label="hello">'
+        '<option selected value="a">Foo</option>'
+        "</optgroup>" in form.a()
+    )
+    assert list(form.a.iter_choices()) == [("a", "Foo", True)]
+
+
+def test_optgroup_shortcut():
+    F = make_form(a=SelectField(choices={"hello": ["foo", "bar"]}))
+    form = F(a="bar")
+
+    assert (
+        '<optgroup label="hello">'
+        '<option value="foo">foo</option>'
+        '<option selected value="bar">bar</option>'
+        "</optgroup>" in form.a()
+    )
+    assert list(form.a.iter_choices()) == [("foo", "foo", False), ("bar", "bar", True)]
+
+
+@pytest.mark.parametrize("choices", [[], ()])
+def test_empty_optgroup(choices):
+    F = make_form(a=SelectField(choices={"hello": choices}))
+    form = F(a="bar")
+    assert '<optgroup label="hello"></optgroup>' in form.a()
+    assert list(form.a.iter_choices()) == []