download improvements

- Add missing showstatus doc
- add name and total to tqdm so the bar is displayed properly
This commit is contained in:
Hellowlol 2017-09-29 23:55:41 +02:00
parent 74eae2acdb
commit 566107697a

View file

@ -242,6 +242,7 @@ def download(url, filename=None, savepath=None, session=None, chunksize=4024,
chunksize (int): What chunksize read/write at the time. chunksize (int): What chunksize read/write at the time.
mocked (bool): Helper to do evertything except write the file. mocked (bool): Helper to do evertything except write the file.
unpack (bool): Unpack the zip file. unpack (bool): Unpack the zip file.
showstatus(bool): Display a progressbar.
Example: Example:
>>> download(a_episode.getStreamURL(), a_episode.location) >>> download(a_episode.getStreamURL(), a_episode.location)
@ -254,10 +255,12 @@ def download(url, filename=None, savepath=None, session=None, chunksize=4024,
# make sure the savepath directory exists # make sure the savepath directory exists
savepath = savepath or os.getcwd() savepath = savepath or os.getcwd()
compat.makedirs(savepath, exist_ok=True) compat.makedirs(savepath, exist_ok=True)
# try getting filename from header if not specified in arguments (used for logs, db) # try getting filename from header if not specified in arguments (used for logs, db)
if not filename and response.headers.get('Content-Disposition'): if not filename and response.headers.get('Content-Disposition'):
filename = re.findall(r'filename=\"(.+)\"', response.headers.get('Content-Disposition')) filename = re.findall(r'filename=\"(.+)\"', response.headers.get('Content-Disposition'))
filename = filename[0] if filename[0] else None filename = filename[0] if filename[0] else None
filename = os.path.basename(filename) filename = os.path.basename(filename)
fullpath = os.path.join(savepath, filename) fullpath = os.path.join(savepath, filename)
# append file.ext from content-type if not already there # append file.ext from content-type if not already there
@ -266,24 +269,31 @@ def download(url, filename=None, savepath=None, session=None, chunksize=4024,
contenttype = response.headers.get('content-type') contenttype = response.headers.get('content-type')
if contenttype and 'image' in contenttype: if contenttype and 'image' in contenttype:
fullpath += contenttype.split('/')[1] fullpath += contenttype.split('/')[1]
# check this is a mocked download (testing) # check this is a mocked download (testing)
if mocked: if mocked:
log.debug('Mocked download %s', fullpath) log.debug('Mocked download %s', fullpath)
return fullpath return fullpath
# save the file to disk # save the file to disk
log.info('Downloading: %s', fullpath) log.info('Downloading: %s', fullpath)
if showstatus: if showstatus:
bar = tqdm(unit='B', unit_scale=True) total = int(response.headers.get('content-length', 0))
bar = tqdm(unit='B', unit_scale=True, total=title, desc=filename)
with open(fullpath, 'wb') as handle: with open(fullpath, 'wb') as handle:
for chunk in response.iter_content(chunk_size=chunksize): for chunk in response.iter_content(chunk_size=chunksize):
handle.write(chunk) handle.write(chunk)
if showstatus: if showstatus:
bar.update(len(chunk)) bar.update(len(chunk))
if showstatus:
bar.close()
# check we want to unzip the contents # check we want to unzip the contents
if fullpath.endswith('zip') and unpack: if fullpath.endswith('zip') and unpack:
with zipfile.ZipFile(fullpath, 'r') as handle: with zipfile.ZipFile(fullpath, 'r') as handle:
handle.extractall(savepath) handle.extractall(savepath)
# finished; return fillpath
return fullpath return fullpath