-
Notifications
You must be signed in to change notification settings - Fork 3
/
utils.py
34 lines (23 loc) · 900 Bytes
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import requests
CHUNK_SIZE = 32768
URL = 'https://docs.google.com/uc?export=download'
def download_file_from_google_drive(id, destination):
print(f'# Downloading {destination}', end=' => ')
session = requests.Session()
response = session.get(URL, params={'id': id}, stream=True)
token = get_confirm_token(response)
if token:
params = {'id': id, 'confirm': token}
response = session.get(URL, params=params, stream=True)
save_response_content(response, destination)
print('Saved')
def get_confirm_token(response):
for key, value in response.cookies.items():
if key.startswith('download_warning'):
return value
return None
def save_response_content(response, destination):
with open(destination, 'wb') as f:
for chunk in response.iter_content(CHUNK_SIZE):
if chunk:
f.write(chunk)