diff --git a/src/flask_wtf/file.py b/src/flask_wtf/file.py index 98c45770..635de929 100644 --- a/src/flask_wtf/file.py +++ b/src/flask_wtf/file.py @@ -1,4 +1,5 @@ from collections import abc +from io import BytesIO from werkzeug.datastructures import FileStorage from wtforms import FileField as _FileField @@ -129,8 +130,17 @@ def __call__(self, form, field): return for f in field_data: - file_size = len(f.read()) - f.seek(0) # reset cursor position to beginning of file + if isinstance(f.stream, BytesIO): + file_size = f.getbuffer().nbytes + elif f.seekable(): + file_size = f.seek(0, 2) + f.seek(0) + else: + file_size = len(f.read()) + try: + f.seek(0) # reset cursor position to beginning of file + except OSError: + pass # welp we broke it if (file_size < self.min_size) or (file_size > self.max_size): # the file is too small or too big => validation failure diff --git a/tests/test_file.py b/tests/test_file.py index 87654d42..00e75183 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -102,6 +102,25 @@ def test_file_size_small_file_passes_validation(form, tmp_path): assert f.validate() +def test_file_size_bytesio_passes_validation(form): + from io import BytesIO + + form.file.kwargs["validators"] = [FileSize(max_size=100)] + contents = BytesIO(b"small") + f = form(file=FileStorage(contents, filename="bytes")) + assert f.validate() + + +def test_file_size_bytesio_too_big_fails_validation(form): + from io import BytesIO + + form.file.kwargs["validators"] = [FileSize(max_size=100)] + contents = BytesIO(b"small" * 30) + f = form(file=FileStorage(contents, filename="bytes")) + assert not f.validate() + assert f.file.errors[0] == "File must be between 0 and 100 bytes." + + @pytest.mark.parametrize( "min_size, max_size, invalid_file_size", [(1, 100, 0), (0, 100, 101)] )