adds the ability to change the label on every object in a queryset for ReferenceFields

Author: sebbekarlsson

flask_mongoengine/wtf/fields.py

@@ -47,13 +47,15 @@ def __init__(
         label_attr="",
         allow_blank=False,
         blank_text="---",
+        label_modifier=None,
         **kwargs,
     ):
 
         super(QuerySetSelectField, self).__init__(label, validators, **kwargs)
         self.label_attr = label_attr
         self.allow_blank = allow_blank
         self.blank_text = blank_text
+        self.label_modifier = label_modifier
         self.queryset = queryset
 
     def iter_choices(self):
@@ -65,7 +67,12 @@ def iter_choices(self):
 
         self.queryset.rewind()
         for obj in self.queryset:
-            label = self.label_attr and getattr(obj, self.label_attr) or obj
+            label = (
+                self.label_modifier(obj)
+                if self.label_modifier
+                else (self.label_attr and getattr(obj, self.label_attr) or obj)
+            )
+
             if isinstance(self.data, list):
                 selected = obj in self.data
             else:

flask_mongoengine/wtf/orm.py

@@ -89,6 +89,11 @@ def convert(self, model, field, field_args):
         if hasattr(field, "to_form_field"):
             return field.to_form_field(model, kwargs)
 
+        if type(field.field) == ReferenceField:
+            kwargs["label_modifier"] = getattr(
+                model, field.name + "_label_modifier", None
+            )
+
         if ftype in self.converters:
             return self.converters[ftype](model, field, kwargs)
 

tests/test_forms.py

@@ -534,6 +534,31 @@ class Post(db.Document):
             form = PostForm()
             self.assertTrue("content-text" in "%s" % form.content.text)
 
+    def test_form_label_modifier(self):
+        with self.app.test_request_context("/"):
+            db = self.db
+
+            class FoodItem(db.Document):
+                title = db.StringField()
+
+            class FoodStore(db.Document):
+                title = db.StringField(max_length=120, required=True)
+                food_items = db.ListField(db.ReferenceField(FoodItem))
+
+                def food_items_label_modifier(obj):
+                    return obj.title
+
+            fruit_names = ["banana", "apple", "pear"]
+
+            food_items = [FoodItem(title=name).save() for name in fruit_names]
+
+            FoodStore(title="John's fruits", food_items=food_items).save()
+
+            FoodStoreForm = model_form(FoodStore)
+            form = FoodStoreForm()
+
+            assert [obj.label.text for obj in form.food_items] == fruit_names
+
 
 if __name__ == "__main__":
     unittest.main()