diff --git a/plexapi/utils.py b/plexapi/utils.py index 38282019..fd007f7c 100644 --- a/plexapi/utils.py +++ b/plexapi/utils.py @@ -242,6 +242,7 @@ def download(url, filename=None, savepath=None, session=None, chunksize=4024, chunksize (int): What chunksize read/write at the time. mocked (bool): Helper to do evertything except write the file. unpack (bool): Unpack the zip file. + showstatus(bool): Display a progressbar. Example: >>> 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 savepath = savepath or os.getcwd() compat.makedirs(savepath, exist_ok=True) + # try getting filename from header if not specified in arguments (used for logs, db) if not filename and response.headers.get('Content-Disposition'): filename = re.findall(r'filename=\"(.+)\"', response.headers.get('Content-Disposition')) filename = filename[0] if filename[0] else None + filename = os.path.basename(filename) fullpath = os.path.join(savepath, filename) # 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') if contenttype and 'image' in contenttype: fullpath += contenttype.split('/')[1] + # check this is a mocked download (testing) if mocked: log.debug('Mocked download %s', fullpath) return fullpath + # save the file to disk log.info('Downloading: %s', fullpath) 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: for chunk in response.iter_content(chunk_size=chunksize): handle.write(chunk) if showstatus: bar.update(len(chunk)) + + if showstatus: + bar.close() # check we want to unzip the contents if fullpath.endswith('zip') and unpack: with zipfile.ZipFile(fullpath, 'r') as handle: handle.extractall(savepath) - # finished; return fillpath + return fullpath