Merge pull request #44 from nagytech/feature/optgroup

Add optgroup support (#43)

Author: azmeuk

wtforms_sqlalchemy/fields.py

@@ -2,6 +2,7 @@
 Useful form fields for use with SQLAlchemy ORM.
 """
 import operator
+from collections import defaultdict
 
 from wtforms import widgets
 from wtforms.fields import SelectFieldBase
@@ -48,6 +49,15 @@ class QuerySelectField(SelectFieldBase):
     model instance and expected to return the label text. Otherwise, the model
     object's `__str__` will be used.
 
+
+    Specify `get_group` to allow `option` elements to be grouped into `optgroup`
+    sections.  If a string, this is the name of an attribute on the model
+    containing the group name.  If a one-argument callable, this callable will
+    be passed the model instance and expected to return a group name.  Otherwise,
+    the `option` elements will not be grouped.  Note: the result of `get_group`
+    will be used as both the grouping key and the display label in the `select`
+    options.
+
     Specify `get_render_kw` to apply HTML attributes to each option. If a
     string, this is the name of an attribute on the model containing a
     dictionary.  If a one-argument callable, this callable will be passed the
@@ -69,6 +79,7 @@ def __init__(
         query_factory=None,
         get_pk=None,
         get_label=None,
+        get_group=None,
         get_render_kw=None,
         allow_blank=False,
         blank_text="",
@@ -93,6 +104,15 @@ def __init__(
         else:
             self.get_label = get_label
 
+        if get_group is None:
+            self._has_groups = False
+        else:
+            self._has_groups = True
+            if isinstance(get_group, str):
+                self.get_group = operator.attrgetter(get_group)
+            else:
+                self.get_group = get_group
+
         if get_render_kw is None:
             self.get_render_kw = lambda _: {}
         elif isinstance(get_render_kw, str):
@@ -133,6 +153,26 @@ def iter_choices(self):
         for pk, obj in self._get_object_list():
             yield (pk, self.get_label(obj), obj == self.data, self.get_render_kw(obj))
 
+    def has_groups(self):
+        return self._has_groups
+
+    def iter_groups(self):
+        if self.has_groups():
+            groups = defaultdict(list)
+            for pk, obj in self._get_object_list():
+                groups[self.get_group(obj)].append((pk, obj))
+            for group, choices in groups.items():
+                yield (group, self._choices_generator(choices))
+
+    def _choices_generator(self, choices):
+        if not choices:
+            _choices = []
+        else:
+            _choices = choices
+
+        for pk, obj in _choices:
+            yield (pk, self.get_label(obj), obj == self.data, self.get_render_kw(obj))
+
     def process_formdata(self, valuelist):
         if valuelist:
             if self.allow_blank and valuelist[0] == "__None":