fields.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192
  1. from mongoengine.base import get_document
  2. from werkzeug.datastructures import FileStorage
  3. from wtforms import fields
  4. try:
  5. from wtforms.fields.core import _unset_value as unset_value
  6. except ImportError:
  7. from wtforms.utils import unset_value
  8. from . import widgets
  9. from flask_admin.model.fields import InlineFormField
  10. def is_empty(file_object):
  11. file_object.seek(0)
  12. first_char = file_object.read(1)
  13. file_object.seek(0)
  14. return not bool(first_char)
  15. class ModelFormField(InlineFormField):
  16. """
  17. Customized ModelFormField for MongoEngine EmbeddedDocuments.
  18. """
  19. def __init__(self, model, view, form_class, form_opts=None, **kwargs):
  20. super(ModelFormField, self).__init__(form_class, **kwargs)
  21. self.model = model
  22. if isinstance(self.model, str):
  23. self.model = get_document(self.model)
  24. self.view = view
  25. self.form_opts = form_opts
  26. def populate_obj(self, obj, name):
  27. candidate = getattr(obj, name, None)
  28. is_created = candidate is None
  29. if is_created:
  30. candidate = self.model()
  31. setattr(obj, name, candidate)
  32. self.form.populate_obj(candidate)
  33. self.view._on_model_change(self.form, candidate, is_created)
  34. class MongoFileField(fields.FileField):
  35. """
  36. GridFS file field.
  37. """
  38. widget = widgets.MongoFileInput()
  39. def __init__(self, label=None, validators=None, **kwargs):
  40. super(MongoFileField, self).__init__(label, validators, **kwargs)
  41. self._should_delete = False
  42. def process(self, formdata, data=unset_value):
  43. if formdata:
  44. marker = '_%s-delete' % self.name
  45. if marker in formdata:
  46. self._should_delete = True
  47. return super(MongoFileField, self).process(formdata, data)
  48. def populate_obj(self, obj, name):
  49. field = getattr(obj, name, None)
  50. if field is not None:
  51. # If field should be deleted, clean it up
  52. if self._should_delete:
  53. field.delete()
  54. return
  55. if isinstance(self.data, FileStorage) and not is_empty(self.data.stream):
  56. if not field.grid_id:
  57. func = field.put
  58. else:
  59. func = field.replace
  60. func(self.data.stream,
  61. filename=self.data.filename,
  62. content_type=self.data.content_type)
  63. class MongoImageField(MongoFileField):
  64. """
  65. GridFS image field.
  66. """
  67. widget = widgets.MongoImageInput()