logo
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions

58 lines
1.5 KiB

import sys
import os
import shutil
import warnings
import requests
import pidfile
from contextlib import contextmanager
from time import sleep
@contextmanager
def exclusive(pidname):
done = False
while not done:
try:
with pidfile.PIDFile(pidname):
yield
done = True
except pidfile.AlreadyRunningError:
sleep(5)
def download_from_url(url, path):
"""Download file, with logic (from tensor2tensor) for Google Drive"""
if 'drive.google.com' not in url:
print('Downloading %s; may take a few minutes' % url)
r = requests.get(url, headers={'User-Agent': 'Mozilla/5.0'})
with open(path, "wb") as file:
file.write(r.content)
return
print('Downloading from Google Drive; may take a few minutes')
confirm_token = None
session = requests.Session()
response = session.get(url, stream=True)
for k, v in response.cookies.items():
if k.startswith("download_warning"):
confirm_token = v
if confirm_token:
url = url + "&confirm=" + confirm_token
response = session.get(url, stream=True)
chunk_size = 16 * 1024
with open(path, "wb") as f:
for chunk in response.iter_content(chunk_size):
if chunk:
f.write(chunk)
class DummyFile(object):
def write(self, x): pass
@contextmanager
def nostdout():
save_stdout = sys.stdout
sys.stdout = DummyFile()
yield
sys.stdout = save_stdout