@@ -157,6 +157,13 @@ def _get_retry_adapter(max_retries: Optional[int] = 5) -> requests.adapters.Base
157
157
return requests .adapters .HTTPAdapter (max_retries = retries )
158
158
159
159
160
+ def get_session (max_retries : Optional [int ] = 5 ) -> requests .Session :
161
+ """Session to hold cookies and connection pool together"""
162
+ session = requests .Session ()
163
+ session .mount ("http" , _get_retry_adapter (max_retries )) # tied to http and https
164
+ return session
165
+
166
+
160
167
def stream_file (
161
168
url : str ,
162
169
fpath : Optional [pathlib .Path ] = None ,
@@ -166,6 +173,7 @@ def stream_file(
166
173
only_first_block : Optional [bool ] = False ,
167
174
max_retries : Optional [int ] = 5 ,
168
175
headers : Optional [Dict [str , str ]] = None ,
176
+ session : Optional [requests .Session ] = None ,
169
177
) -> Union [int , requests .structures .CaseInsensitiveDict ]:
170
178
"""Stream data from a URL to either a BytesIO object or a file
171
179
Arguments -
@@ -175,16 +183,17 @@ def stream_file(
175
183
proxies - A dict of proxies to be used
176
184
https://requests.readthedocs.io/en/master/user/advanced/#proxies
177
185
only_first_block - Whether to download only one (first) block
178
- max_retries - Maximum number of retries after which error is raised
186
+ max_retries - Maximum number of retries after which error is raised. Does not
187
+ apply if using your own session
188
+ session - Session object to make the request with. A new one created otherwise
179
189
Returns the total number of bytes downloaded and the response headers"""
180
190
181
191
# if no output option is supplied
182
192
if fpath is None and byte_stream is None :
183
193
raise ValueError ("Either file path or a bytesIO object is needed" )
184
194
185
- session = requests .Session ()
186
- retry_adapter = _get_retry_adapter (max_retries )
187
- session .mount ("http" , retry_adapter ) # tied to http and https
195
+ if not session :
196
+ session = get_session (max_retries )
188
197
resp = session .get (
189
198
url ,
190
199
stream = True ,
0 commit comments