file.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. from collections import abc
  2. from werkzeug.datastructures import FileStorage
  3. from wtforms import FileField as _FileField
  4. from wtforms import MultipleFileField as _MultipleFileField
  5. from wtforms.validators import DataRequired
  6. from wtforms.validators import StopValidation
  7. from wtforms.validators import ValidationError
  8. class FileField(_FileField):
  9. """Werkzeug-aware subclass of :class:`wtforms.fields.FileField`."""
  10. def process_formdata(self, valuelist):
  11. valuelist = (x for x in valuelist if isinstance(x, FileStorage) and x)
  12. data = next(valuelist, None)
  13. if data is not None:
  14. self.data = data
  15. else:
  16. self.raw_data = ()
  17. class MultipleFileField(_MultipleFileField):
  18. """Werkzeug-aware subclass of :class:`wtforms.fields.MultipleFileField`.
  19. .. versionadded:: 1.2.0
  20. """
  21. def process_formdata(self, valuelist):
  22. valuelist = (x for x in valuelist if isinstance(x, FileStorage) and x)
  23. data = list(valuelist) or None
  24. if data is not None:
  25. self.data = data
  26. else:
  27. self.raw_data = ()
  28. class FileRequired(DataRequired):
  29. """Validates that the uploaded files(s) is a Werkzeug
  30. :class:`~werkzeug.datastructures.FileStorage` object.
  31. :param message: error message
  32. You can also use the synonym ``file_required``.
  33. """
  34. def __call__(self, form, field):
  35. field_data = [field.data] if not isinstance(field.data, list) else field.data
  36. if not (
  37. all(isinstance(x, FileStorage) and x for x in field_data) and field_data
  38. ):
  39. raise StopValidation(
  40. self.message or field.gettext("This field is required.")
  41. )
  42. file_required = FileRequired
  43. class FileAllowed:
  44. """Validates that the uploaded file(s) is allowed by a given list of
  45. extensions or a Flask-Uploads :class:`~flaskext.uploads.UploadSet`.
  46. :param upload_set: A list of extensions or an
  47. :class:`~flaskext.uploads.UploadSet`
  48. :param message: error message
  49. You can also use the synonym ``file_allowed``.
  50. """
  51. def __init__(self, upload_set, message=None):
  52. self.upload_set = upload_set
  53. self.message = message
  54. def __call__(self, form, field):
  55. field_data = [field.data] if not isinstance(field.data, list) else field.data
  56. if not (
  57. all(isinstance(x, FileStorage) and x for x in field_data) and field_data
  58. ):
  59. return
  60. filenames = [f.filename.lower() for f in field_data]
  61. for filename in filenames:
  62. if isinstance(self.upload_set, abc.Iterable):
  63. if any(filename.endswith("." + x) for x in self.upload_set):
  64. continue
  65. raise StopValidation(
  66. self.message
  67. or field.gettext(
  68. "File does not have an approved extension: {extensions}"
  69. ).format(extensions=", ".join(self.upload_set))
  70. )
  71. if not self.upload_set.file_allowed(field_data, filename):
  72. raise StopValidation(
  73. self.message
  74. or field.gettext("File does not have an approved extension.")
  75. )
  76. file_allowed = FileAllowed
  77. class FileSize:
  78. """Validates that the uploaded file(s) is within a minimum and maximum
  79. file size (set in bytes).
  80. :param min_size: minimum allowed file size (in bytes). Defaults to 0 bytes.
  81. :param max_size: maximum allowed file size (in bytes).
  82. :param message: error message
  83. You can also use the synonym ``file_size``.
  84. """
  85. def __init__(self, max_size, min_size=0, message=None):
  86. self.min_size = min_size
  87. self.max_size = max_size
  88. self.message = message
  89. def __call__(self, form, field):
  90. field_data = [field.data] if not isinstance(field.data, list) else field.data
  91. if not (
  92. all(isinstance(x, FileStorage) and x for x in field_data) and field_data
  93. ):
  94. return
  95. for f in field_data:
  96. file_size = len(f.read())
  97. f.seek(0) # reset cursor position to beginning of file
  98. if (file_size < self.min_size) or (file_size > self.max_size):
  99. # the file is too small or too big => validation failure
  100. raise ValidationError(
  101. self.message
  102. or field.gettext(
  103. "File must be between {min_size} and {max_size} bytes.".format(
  104. min_size=self.min_size, max_size=self.max_size
  105. )
  106. )
  107. )
  108. file_size = FileSize