Mercurial > repos > shellac > sam_consensus_v3
comparison env/lib/python3.9/site-packages/requests_toolbelt/downloadutils/tee.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 """Tee function implementations.""" | |
2 import io | |
3 | |
4 _DEFAULT_CHUNKSIZE = 65536 | |
5 | |
6 __all__ = ['tee', 'tee_to_file', 'tee_to_bytearray'] | |
7 | |
8 | |
9 def _tee(response, callback, chunksize, decode_content): | |
10 for chunk in response.raw.stream(amt=chunksize, | |
11 decode_content=decode_content): | |
12 callback(chunk) | |
13 yield chunk | |
14 | |
15 | |
16 def tee(response, fileobject, chunksize=_DEFAULT_CHUNKSIZE, | |
17 decode_content=None): | |
18 """Stream the response both to the generator and a file. | |
19 | |
20 This will stream the response body while writing the bytes to | |
21 ``fileobject``. | |
22 | |
23 Example usage: | |
24 | |
25 .. code-block:: python | |
26 | |
27 resp = requests.get(url, stream=True) | |
28 with open('save_file', 'wb') as save_file: | |
29 for chunk in tee(resp, save_file): | |
30 # do stuff with chunk | |
31 | |
32 .. code-block:: python | |
33 | |
34 import io | |
35 | |
36 resp = requests.get(url, stream=True) | |
37 fileobject = io.BytesIO() | |
38 | |
39 for chunk in tee(resp, fileobject): | |
40 # do stuff with chunk | |
41 | |
42 :param response: Response from requests. | |
43 :type response: requests.Response | |
44 :param fileobject: Writable file-like object. | |
45 :type fileobject: file, io.BytesIO | |
46 :param int chunksize: (optional), Size of chunk to attempt to stream. | |
47 :param bool decode_content: (optional), If True, this will decode the | |
48 compressed content of the response. | |
49 :raises: TypeError if the fileobject wasn't opened with the right mode | |
50 or isn't a BytesIO object. | |
51 """ | |
52 # We will be streaming the raw bytes from over the wire, so we need to | |
53 # ensure that writing to the fileobject will preserve those bytes. On | |
54 # Python3, if the user passes an io.StringIO, this will fail, so we need | |
55 # to check for BytesIO instead. | |
56 if not ('b' in getattr(fileobject, 'mode', '') or | |
57 isinstance(fileobject, io.BytesIO)): | |
58 raise TypeError('tee() will write bytes directly to this fileobject' | |
59 ', it must be opened with the "b" flag if it is a file' | |
60 ' or inherit from io.BytesIO.') | |
61 | |
62 return _tee(response, fileobject.write, chunksize, decode_content) | |
63 | |
64 | |
65 def tee_to_file(response, filename, chunksize=_DEFAULT_CHUNKSIZE, | |
66 decode_content=None): | |
67 """Stream the response both to the generator and a file. | |
68 | |
69 This will open a file named ``filename`` and stream the response body | |
70 while writing the bytes to the opened file object. | |
71 | |
72 Example usage: | |
73 | |
74 .. code-block:: python | |
75 | |
76 resp = requests.get(url, stream=True) | |
77 for chunk in tee_to_file(resp, 'save_file'): | |
78 # do stuff with chunk | |
79 | |
80 :param response: Response from requests. | |
81 :type response: requests.Response | |
82 :param str filename: Name of file in which we write the response content. | |
83 :param int chunksize: (optional), Size of chunk to attempt to stream. | |
84 :param bool decode_content: (optional), If True, this will decode the | |
85 compressed content of the response. | |
86 """ | |
87 with open(filename, 'wb') as fd: | |
88 for chunk in tee(response, fd, chunksize, decode_content): | |
89 yield chunk | |
90 | |
91 | |
92 def tee_to_bytearray(response, bytearr, chunksize=_DEFAULT_CHUNKSIZE, | |
93 decode_content=None): | |
94 """Stream the response both to the generator and a bytearray. | |
95 | |
96 This will stream the response provided to the function, add them to the | |
97 provided :class:`bytearray` and yield them to the user. | |
98 | |
99 .. note:: | |
100 | |
101 This uses the :meth:`bytearray.extend` by default instead of passing | |
102 the bytearray into the ``readinto`` method. | |
103 | |
104 Example usage: | |
105 | |
106 .. code-block:: python | |
107 | |
108 b = bytearray() | |
109 resp = requests.get(url, stream=True) | |
110 for chunk in tee_to_bytearray(resp, b): | |
111 # do stuff with chunk | |
112 | |
113 :param response: Response from requests. | |
114 :type response: requests.Response | |
115 :param bytearray bytearr: Array to add the streamed bytes to. | |
116 :param int chunksize: (optional), Size of chunk to attempt to stream. | |
117 :param bool decode_content: (optional), If True, this will decode the | |
118 compressed content of the response. | |
119 """ | |
120 if not isinstance(bytearr, bytearray): | |
121 raise TypeError('tee_to_bytearray() expects bytearr to be a ' | |
122 'bytearray') | |
123 return _tee(response, bytearr.extend, chunksize, decode_content) |