[5ec3a0b] | 1 | from shutil import copyfileobj |
---|
| 2 | |
---|
| 3 | from pyramid.i18n import TranslationStringFactory |
---|
| 4 | |
---|
| 5 | from formencode.validators import FieldStorageUploadConverter, Invalid |
---|
| 6 | from ZODB.blob import Blob |
---|
| 7 | |
---|
| 8 | |
---|
| 9 | _ = TranslationStringFactory('OpenWorkouts') |
---|
| 10 | |
---|
| 11 | |
---|
| 12 | class FieldStorageBlob(FieldStorageUploadConverter): |
---|
| 13 | messages = dict( |
---|
| 14 | badExtension=_('Please upload a valid file type.')) |
---|
| 15 | |
---|
| 16 | def __init__(self, *args, **kw): |
---|
| 17 | self.whitelist = kw.get('whitelist', None) |
---|
| 18 | super(FieldStorageBlob, self).__init__(*args, **kw) |
---|
| 19 | |
---|
| 20 | def _convert_to_python(self, value, state): |
---|
| 21 | """ |
---|
| 22 | Converts a field upload into a ZODB blob |
---|
| 23 | |
---|
| 24 | >>> from io import BytesIO |
---|
| 25 | >>> from contextlib import contextmanager |
---|
| 26 | >>> ifile = lambda: 1 |
---|
| 27 | >>> @contextmanager |
---|
| 28 | ... def likefile(): |
---|
| 29 | ... yield BytesIO(b'1234') |
---|
| 30 | >>> ifile.file = likefile() |
---|
| 31 | >>> ifile.filename = 'test.pdf' |
---|
| 32 | >>> validator = FieldStorageBlob() |
---|
| 33 | >>> blob = validator._convert_to_python(ifile, None) |
---|
| 34 | >>> blob # doctest: +ELLIPSIS |
---|
| 35 | <ZODB.blob.Blob object at 0x...> |
---|
| 36 | >>> blob.open('r').read() |
---|
| 37 | b'1234' |
---|
| 38 | >>> blob.file_extension |
---|
| 39 | 'pdf' |
---|
| 40 | """ |
---|
| 41 | value = super(FieldStorageBlob, self)._convert_to_python(value, state) |
---|
| 42 | file_extension = value.filename.rsplit('.', 1)[-1] |
---|
| 43 | blob = Blob() |
---|
| 44 | blob.file_extension = file_extension |
---|
| 45 | with value.file as infile, blob.open('w') as out: |
---|
| 46 | infile.seek(0) |
---|
| 47 | copyfileobj(infile, out) |
---|
| 48 | return blob |
---|
| 49 | |
---|
| 50 | def _validate_python(self, value, state): |
---|
| 51 | """ |
---|
| 52 | Validates the uploaded file versus a whitelist of file extensions |
---|
| 53 | |
---|
| 54 | >>> from unittest.mock import Mock |
---|
| 55 | >>> validator = FieldStorageBlob() |
---|
| 56 | >>> value = Mock(file_extension='bla') |
---|
| 57 | >>> validator._validate_python(value, None) |
---|
| 58 | |
---|
| 59 | >>> validator2 = FieldStorageBlob(whitelist=['bla', 'jpeg']) |
---|
| 60 | >>> validator2._validate_python(value, None) |
---|
| 61 | |
---|
| 62 | >>> validator3 = FieldStorageBlob(whitelist=['pdf', 'jpeg']) |
---|
| 63 | >>> validator3._validate_python(value, None) |
---|
| 64 | Traceback (most recent call last): |
---|
| 65 | ... |
---|
| 66 | formencode.api.Invalid: Please upload a valid file type. |
---|
| 67 | """ |
---|
| 68 | if self.whitelist is None: |
---|
| 69 | return |
---|
| 70 | |
---|
| 71 | if value.file_extension not in self.whitelist: |
---|
| 72 | raise Invalid(self.message('badExtension', state), value, state) |
---|