Mercurial > repos > shellac > sam_consensus_v3
comparison env/lib/python3.9/site-packages/requests_toolbelt/downloadutils/stream.py @ 0:4f3585e2f14b draft default tip
"planemo upload commit 60cee0fc7c0cda8592644e1aad72851dec82c959"
author | shellac |
---|---|
date | Mon, 22 Mar 2021 18:12:50 +0000 |
parents | |
children |
comparison
equal
deleted
inserted
replaced
-1:000000000000 | 0:4f3585e2f14b |
---|---|
1 # -*- coding: utf-8 -*- | |
2 """Utilities for dealing with streamed requests.""" | |
3 import collections | |
4 import os.path | |
5 import re | |
6 | |
7 from .. import exceptions as exc | |
8 | |
9 # Regular expressions stolen from werkzeug/http.py | |
10 # cd2c97bb0a076da2322f11adce0b2731f9193396 L62-L64 | |
11 _QUOTED_STRING_RE = r'"[^"\\]*(?:\\.[^"\\]*)*"' | |
12 _OPTION_HEADER_PIECE_RE = re.compile( | |
13 r';\s*(%s|[^\s;=]+)\s*(?:=\s*(%s|[^;]+))?\s*' % (_QUOTED_STRING_RE, | |
14 _QUOTED_STRING_RE) | |
15 ) | |
16 _DEFAULT_CHUNKSIZE = 512 | |
17 | |
18 | |
19 def _get_filename(content_disposition): | |
20 for match in _OPTION_HEADER_PIECE_RE.finditer(content_disposition): | |
21 k, v = match.groups() | |
22 if k == 'filename': | |
23 # ignore any directory paths in the filename | |
24 return os.path.split(v)[1] | |
25 return None | |
26 | |
27 | |
28 def get_download_file_path(response, path): | |
29 """ | |
30 Given a response and a path, return a file path for a download. | |
31 | |
32 If a ``path`` parameter is a directory, this function will parse the | |
33 ``Content-Disposition`` header on the response to determine the name of the | |
34 file as reported by the server, and return a file path in the specified | |
35 directory. | |
36 | |
37 If ``path`` is empty or None, this function will return a path relative | |
38 to the process' current working directory. | |
39 | |
40 If path is a full file path, return it. | |
41 | |
42 :param response: A Response object from requests | |
43 :type response: requests.models.Response | |
44 :param str path: Directory or file path. | |
45 :returns: full file path to download as | |
46 :rtype: str | |
47 :raises: :class:`requests_toolbelt.exceptions.StreamingError` | |
48 """ | |
49 path_is_dir = path and os.path.isdir(path) | |
50 | |
51 if path and not path_is_dir: | |
52 # fully qualified file path | |
53 filepath = path | |
54 else: | |
55 response_filename = _get_filename( | |
56 response.headers.get('content-disposition', '') | |
57 ) | |
58 if not response_filename: | |
59 raise exc.StreamingError('No filename given to stream response to') | |
60 | |
61 if path_is_dir: | |
62 # directory to download to | |
63 filepath = os.path.join(path, response_filename) | |
64 else: | |
65 # fallback to downloading to current working directory | |
66 filepath = response_filename | |
67 | |
68 return filepath | |
69 | |
70 | |
71 def stream_response_to_file(response, path=None, chunksize=_DEFAULT_CHUNKSIZE): | |
72 """Stream a response body to the specified file. | |
73 | |
74 Either use the ``path`` provided or use the name provided in the | |
75 ``Content-Disposition`` header. | |
76 | |
77 .. warning:: | |
78 | |
79 If you pass this function an open file-like object as the ``path`` | |
80 parameter, the function will not close that file for you. | |
81 | |
82 .. warning:: | |
83 | |
84 This function will not automatically close the response object | |
85 passed in as the ``response`` parameter. | |
86 | |
87 If a ``path`` parameter is a directory, this function will parse the | |
88 ``Content-Disposition`` header on the response to determine the name of the | |
89 file as reported by the server, and return a file path in the specified | |
90 directory. If no ``path`` parameter is supplied, this function will default | |
91 to the process' current working directory. | |
92 | |
93 .. code-block:: python | |
94 | |
95 import requests | |
96 from requests_toolbelt import exceptions | |
97 from requests_toolbelt.downloadutils import stream | |
98 | |
99 r = requests.get(url, stream=True) | |
100 try: | |
101 filename = stream.stream_response_to_file(r) | |
102 except exceptions.StreamingError as e: | |
103 # The toolbelt could not find the filename in the | |
104 # Content-Disposition | |
105 print(e.message) | |
106 | |
107 You can also specify the filename as a string. This will be passed to | |
108 the built-in :func:`open` and we will read the content into the file. | |
109 | |
110 .. code-block:: python | |
111 | |
112 import requests | |
113 from requests_toolbelt.downloadutils import stream | |
114 | |
115 r = requests.get(url, stream=True) | |
116 filename = stream.stream_response_to_file(r, path='myfile') | |
117 | |
118 If the calculated download file path already exists, this function will | |
119 raise a StreamingError. | |
120 | |
121 Instead, if you want to manage the file object yourself, you need to | |
122 provide either a :class:`io.BytesIO` object or a file opened with the | |
123 `'b'` flag. See the two examples below for more details. | |
124 | |
125 .. code-block:: python | |
126 | |
127 import requests | |
128 from requests_toolbelt.downloadutils import stream | |
129 | |
130 with open('myfile', 'wb') as fd: | |
131 r = requests.get(url, stream=True) | |
132 filename = stream.stream_response_to_file(r, path=fd) | |
133 | |
134 print('{0} saved to {1}'.format(url, filename)) | |
135 | |
136 .. code-block:: python | |
137 | |
138 import io | |
139 import requests | |
140 from requests_toolbelt.downloadutils import stream | |
141 | |
142 b = io.BytesIO() | |
143 r = requests.get(url, stream=True) | |
144 filename = stream.stream_response_to_file(r, path=b) | |
145 assert filename is None | |
146 | |
147 :param response: A Response object from requests | |
148 :type response: requests.models.Response | |
149 :param path: *(optional)*, Either a string with the path to the location | |
150 to save the response content, or a file-like object expecting bytes. | |
151 :type path: :class:`str`, or object with a :meth:`write` | |
152 :param int chunksize: (optional), Size of chunk to attempt to stream | |
153 (default 512B). | |
154 :returns: The name of the file, if one can be determined, else None | |
155 :rtype: str | |
156 :raises: :class:`requests_toolbelt.exceptions.StreamingError` | |
157 """ | |
158 pre_opened = False | |
159 fd = None | |
160 filename = None | |
161 if path and isinstance(getattr(path, 'write', None), collections.Callable): | |
162 pre_opened = True | |
163 fd = path | |
164 filename = getattr(fd, 'name', None) | |
165 else: | |
166 filename = get_download_file_path(response, path) | |
167 if os.path.exists(filename): | |
168 raise exc.StreamingError("File already exists: %s" % filename) | |
169 fd = open(filename, 'wb') | |
170 | |
171 for chunk in response.iter_content(chunk_size=chunksize): | |
172 fd.write(chunk) | |
173 | |
174 if not pre_opened: | |
175 fd.close() | |
176 | |
177 return filename |