1# 

2# Copyright (C) 2019-2020 Leo P. Singer <leo.singer@ligo.org> 

3# 

4# This program is free software: you can redistribute it and/or modify 

5# it under the terms of the GNU General Public License as published by 

6# the Free Software Foundation, either version 3 of the License, or 

7# (at your option) any later version. 

8# 

9# This program is distributed in the hope that it will be useful, 

10# but WITHOUT ANY WARRANTY; without even the implied warranty of 

11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

12# GNU General Public License for more details. 

13# 

14# You should have received a copy of the GNU General Public License 

15# along with this program. If not, see <https://www.gnu.org/licenses/>. 

16# 

17from __future__ import absolute_import 

18from os.path import basename 

19from mimetypes import guess_type 

20from sys import stdin 

21 

22from requests.utils import guess_filename, to_key_val_list 

23 

24 

25def _guess_content_type(filename): 

26 return guess_type(filename)[0] or 'application/octet-stream' 

27 

28 

29def _guess_mime_type(key, val): 

30 if not isinstance(val, (tuple, list)): 

31 filename = guess_filename(val) or key 

32 filetype = _guess_content_type(filename) 

33 val = (filename, val, filetype) 

34 elif len(val) < 3 or val[2] is None: 

35 filename = val[0] 

36 filetype = _guess_content_type(filename) 

37 # val = (*val[:2], filetype, *val[3:]) 

38 # FIXME: Python 2 

39 val = tuple(val[:2]) + (filetype,) + tuple(val[3:]) 

40 return key, val 

41 

42 

43def _read_files(key, val): 

44 if isinstance(val, (tuple, list)) and (len(val) < 2 or val[1] is None): 

45 filename = val[0] 

46 with (stdin.buffer if filename == '-' else open(filename, 'rb')) as f: 

47 data = f.read() 

48 filename = basename(filename) 

49 # val = (filename, data, *val[2:]) 

50 # FIXME: Python 2 

51 val = (filename, data) + tuple(val[2:]) 

52 return key, val 

53 

54 

55def _prepare_file(key, val): 

56 key, val = _read_files(key, val) 

57 key, val = _guess_mime_type(key, val) 

58 return key, val 

59 

60 

61def _prepare_files(files): 

62 if files is not None: 

63 files = [_prepare_file(k, v) for k, v in to_key_val_list(files)] 

64 return files 

65 

66 

67class SessionFileMixin(object): 

68 """A mixin for :class:`requests.Session` to add features for file uploads. 

69 

70 The :meth:`requests.Session.request` method takes a `files` argument which 

71 is a dictionary of `{fieldname: fileobject}`, where `fileobject` may be a 

72 file-like object or a tuple of 2-4 elements consisting of the filename, 

73 file content, MIME type, and any custom headers. 

74 

75 This mixin adds the following features: 

76 

77 * The MIME type is automatically guessed from the filename. 

78 * If the file content is None, then the filename is treated as a path, and 

79 the file is opened and read. If the filename is `-`, then the file 

80 content is read from stdin. 

81 """ 

82 

83 def request( 

84 self, method, url, params=None, data=None, headers=None, 

85 cookies=None, files=None, auth=None, timeout=None, 

86 allow_redirects=True, proxies=None, hooks=None, stream=None, 

87 verify=None, cert=None, json=None): 

88 return super(SessionFileMixin, self).request( 

89 method, url, params=params, data=data, headers=headers, 

90 cookies=cookies, files=_prepare_files(files), auth=auth, 

91 timeout=timeout, allow_redirects=True, proxies=proxies, 

92 hooks=hooks, stream=stream, verify=verify, cert=cert, json=json)