QuerySetSelectField and QuerySetSelectMultipleField improvements

- code cleanup
- remove useless use of clone() (MongoEngine/mongoengine@9bbd8db)
- filter by pk not id
- fix QuerySetSelectMultipleField to allow emptying the value list (fixes #148)
- added a few tests

Author: lafrech

flask_mongoengine/wtf/fields.py

@@ -37,8 +37,10 @@ class QuerySetSelectField(SelectFieldBase):
     """
     widget = widgets.Select()
 
-    def __init__(self, label=u'', validators=None, queryset=None, label_attr='',
-                 allow_blank=False, blank_text=u'---', **kwargs):
+    def __init__(self, label=u'', validators=None, queryset=None,
+                 label_attr='', allow_blank=False, blank_text=u'---',
+                 **kwargs):
+
         super(QuerySetSelectField, self).__init__(label, validators, **kwargs)
         self.label_attr = label_attr
         self.allow_blank = allow_blank
@@ -71,8 +73,7 @@ def process_formdata(self, valuelist):
                     return
 
                 try:
-                    # clone() because of https://github.com/MongoEngine/mongoengine/issues/56
-                    obj = self.queryset.clone().get(pk=valuelist[0])
+                    obj = self.queryset.get(pk=valuelist[0])
                     self.data = obj
                 except DoesNotExist:
                     self.data = None
@@ -90,11 +91,16 @@ class QuerySetSelectMultipleField(QuerySetSelectField):
 
     widget = widgets.Select(multiple=True)
 
-    def  __init__(self, label=u'', validators=None, queryset=None, label_attr='',
-                  allow_blank=False, blank_text=u'---', **kwargs):
-        super(QuerySetSelectMultipleField, self).__init__(label, validators, queryset, label_attr, allow_blank, blank_text, **kwargs)
+    def __init__(self, label=u'', validators=None, queryset=None,
+                 label_attr='', allow_blank=False, blank_text=u'---',
+                 **kwargs):
+
+        super(QuerySetSelectMultipleField, self).__init__(
+            label, validators, queryset, label_attr, allow_blank, blank_text,
+            **kwargs)
 
     def process_formdata(self, valuelist):
+
         if valuelist:
             if valuelist[0] == '__None':
                 self.data = None
@@ -104,10 +110,14 @@ def process_formdata(self, valuelist):
                     return
 
                 self.queryset.rewind()
-                self.data = [obj for obj in self.queryset if str(obj.id) in valuelist]
+                self.data = list(self.queryset(pk__in=valuelist))
                 if not len(self.data):
                     self.data = None
 
+        # If no value passed, empty the list
+        else:
+            self.data = None
+
     def _is_selected(self, item):
         return item in self.data if self.data else False
 

tests/test_forms.py

@@ -244,6 +244,18 @@ class DogOwner(db.Document):
             self.assertEqual(wtforms.widgets.Select, type(form.dog.widget))
             self.assertFalse(form.dog.widget.multiple)
 
+            # Validate selecting one item
+            form = DogOwnerForm(MultiDict({
+                'dog': dog.id,
+                }))
+            self.assertEqual(form.dog.data, dog)
+
+            # Validate selecting no item
+            form = DogOwnerForm(MultiDict({
+                'dog': u'__None',
+                }), dog=dog)
+            self.assertEqual(form.dog.data, None)
+
     def test_modelselectfield_multiple(self):
         with self.app.test_request_context('/'):
             db = self.db
@@ -272,6 +284,18 @@ class DogOwner(db.Document):
             self.assertTrue(choices[0].checked)
             self.assertTrue(choices[1].checked)
 
+            # Validate selecting two items
+            form = DogOwnerForm(MultiDict({
+                'dogs': [dog.id for dog in dogs],
+                }))
+            self.assertEqual(form.dogs.data, dogs)
+
+            # Validate selecting none actually empties the list
+            form = DogOwnerForm(MultiDict({
+                'dogs': [],
+                }), dogs=dogs)
+            self.assertEqual(form.dogs.data, None)
+
     def test_modelselectfield_multiple_initalvalue_None(self):
         with self.app.test_request_context('/'):
             db = self.db