Fixed only / exclude in model forms (#49)

Author: rozza

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 
 Changes in 0.7
 ==============
+- Fixed only / exclude in model forms (#49)
 - Added automatic choices coerce for simple types (#34)
 - Fixed EmailField and URLField rendering and validation (#44, #9)
 - Use help_text for field description (#43)

flask_mongoengine/wtf/orm.py

@@ -226,10 +226,9 @@ def model_fields(model, only=None, exclude=None, field_args=None, converter=None
     field_names = map(itemgetter(0), sorted(names, key=itemgetter(1)))
 
     if only:
-        field_names = set(field_names)
-        field_names = (x for x in only if x in field_names)
+        field_names = set((x for x in only if x in set(field_names)))
     elif exclude:
-        field_names = (x for x in field_names if x not in exclude)
+        field_names = set((x for x in set(field_names) if x not in exclude))
 
     field_dict = OrderedDict()
     for name in field_names:

tests/test_forms.py

@@ -150,6 +150,27 @@ class LinkPost(BlogPost):
 
             self.assertEqual(post.tags, ['flask', 'mongodb', 'mongoengine', 'flask-mongoengine'])
 
+    def test_model_form_only(self):
+        with self.app.test_request_context('/'):
+            db = self.db
+
+            class BlogPost(db.Document):
+                title = db.StringField(required=True, max_length=200)
+                posted = db.DateTimeField(default=datetime.datetime.now)
+                tags = db.ListField(db.StringField())
+
+            BlogPost.drop_collection()
+
+            BlogPostForm = model_form(BlogPost, only=['tags'])
+            form = BlogPostForm()
+            self.assertTrue(hasattr(form, 'tags'))
+            self.assertFalse(hasattr(form, 'posted'))
+
+            BlogPostForm = model_form(BlogPost, exclude=['posted'])
+            form = BlogPostForm()
+            self.assertTrue(hasattr(form, 'tags'))
+            self.assertFalse(hasattr(form, 'posted'))
+
     def test_model_form_with_custom_query_set(self):
         with self.app.test_request_context('/'):
             db = self.db
@@ -335,5 +356,6 @@ class Post(db.Document):
             self.assertTrue("content-text" in "%s" % form.content.text)
 
 
+
 if __name__ == '__main__':
     unittest.main()