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) |
---|