voussoir/voussoirkit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Use DotDict for the download plans.

Browse source
This commit is contained in:
Ethan Dalool 2020-09-25 15:56:24 -07:00
parent 0948139883
commit 0613fd10e4
1 changed files with 29 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
voussoirkit/downloady.py
View file

@ -7,6 +7,7 @@ import warnings
from voussoirkit import bytestring from voussoirkit import bytestring
from voussoirkit import clipext from voussoirkit import clipext
from voussoirkit import dotdict
from voussoirkit import pathclass from voussoirkit import pathclass
from voussoirkit import ratelimiter from voussoirkit import ratelimiter
from voussoirkit import safeprint from voussoirkit import safeprint
@ -77,14 +78,15 @@ def download_file(
timeout=timeout, timeout=timeout,
verify_ssl=verify_ssl, verify_ssl=verify_ssl,
) )
if plan is None: if plan is None:
return return
return download_plan(plan) return download_plan(plan)
def download_plan(plan): def download_plan(plan):
temp_localname = plan['download_into'] temp_localname = plan.download_into
real_localname = plan['real_localname'] real_localname = plan.real_localname
directory = os.path.split(temp_localname)[0] directory = os.path.split(temp_localname)[0]
if directory != '': if directory != '':
os.makedirs(directory, exist_ok=True) os.makedirs(directory, exist_ok=True)
@ -92,38 +94,38 @@ def download_plan(plan):
if not is_special_file(temp_localname): if not is_special_file(temp_localname):
touch(temp_localname) touch(temp_localname)
if plan['plan_type'] in ['resume', 'partial']: if plan.plan_type in ['resume', 'partial']:
file_handle = open(temp_localname, 'r+b') file_handle = open(temp_localname, 'r+b')
file_handle.seek(plan['seek_to']) file_handle.seek(plan.seek_to)
bytes_downloaded = plan['seek_to'] bytes_downloaded = plan.seek_to
elif plan['plan_type'] == 'fulldownload': elif plan.plan_type == 'fulldownload':
file_handle = open(temp_localname, 'wb') file_handle = open(temp_localname, 'wb')
bytes_downloaded = 0 bytes_downloaded = 0
if plan['header_range_min'] is not None: if plan.header_range_min is not None:
plan['headers']['range'] = 'bytes={min}-{max}'.format( plan.headers['range'] = 'bytes={min}-{max}'.format(
min=plan['header_range_min'], min=plan.header_range_min,
max=plan['header_range_max'], max=plan.header_range_max,
) )
download_stream = request( download_stream = request(
'get', 'get',
plan['url'], plan.url,
stream=True, stream=True,
auth=plan['auth'], auth=plan.auth,
headers=plan['headers'], headers=plan.headers,
timeout=plan['timeout'], timeout=plan.timeout,
verify_ssl=plan['verify_ssl'], verify_ssl=plan.verify_ssl,
) )
if plan['remote_total_bytes'] is None: if plan.remote_total_bytes is None:
# Since we didn't do a head, let's fill this in now. # Since we didn't do a head, let's fill this in now.
plan['remote_total_bytes'] = int(download_stream.headers.get('Content-Length', 0)) plan.remote_total_bytes = int(download_stream.headers.get('Content-Length', 0))
callback_progress = plan['callback_progress'] callback_progress = plan.callback_progress
if callback_progress is not None: if callback_progress is not None:
callback_progress = callback_progress(plan['remote_total_bytes']) callback_progress = callback_progress(plan.remote_total_bytes)
for chunk in download_stream.iter_content(chunk_size=CHUNKSIZE): for chunk in download_stream.iter_content(chunk_size=CHUNKSIZE):
bytes_downloaded += len(chunk) bytes_downloaded += len(chunk)
@ -131,18 +133,18 @@ def download_plan(plan):
if callback_progress is not None: if callback_progress is not None:
callback_progress.step(bytes_downloaded) callback_progress.step(bytes_downloaded)
if plan['limiter'] is not None and bytes_downloaded < plan['remote_total_bytes']: if plan.limiter is not None and bytes_downloaded < plan.remote_total_bytes:
plan['limiter'].limit(len(chunk)) plan.limiter.limit(len(chunk))
file_handle.close() file_handle.close()
# Don't try to rename /dev/null or other special names # Don't try to rename /dev/null or other special names
if not is_special_file(temp_localname) and not is_special_file(real_localname): if not is_special_file(temp_localname) and not is_special_file(real_localname):
localsize = os.path.getsize(temp_localname) localsize = os.path.getsize(temp_localname)
undersized = plan['plan_type'] != 'partial' and localsize < plan['remote_total_bytes'] undersized = plan.plan_type != 'partial' and localsize < plan.remote_total_bytes
if plan['raise_for_undersized'] and undersized: if plan.raise_for_undersized and undersized:
message = 'File does not contain expected number of bytes. Received {size} / {total}' message = 'File does not contain expected number of bytes. Received {size} / {total}'
message = message.format(size=localsize, total=plan['remote_total_bytes']) message = message.format(size=localsize, total=plan.remote_total_bytes)
raise NotEnoughBytes(message) raise NotEnoughBytes(message)
if temp_localname != real_localname: if temp_localname != real_localname:
@ -235,7 +237,7 @@ def prepare_plan(
'timeout': timeout, 'timeout': timeout,
'verify_ssl': verify_ssl, 'verify_ssl': verify_ssl,
} }
plan_fulldownload = dict( plan_fulldownload = dotdict.DotDict(
plan_base, plan_base,
download_into=temp_localname, download_into=temp_localname,
header_range_min=None, header_range_min=None,
@ -243,7 +245,7 @@ def prepare_plan(
plan_type='fulldownload', plan_type='fulldownload',
seek_to=0, seek_to=0,
) )
plan_resume = dict( plan_resume = dotdict.DotDict(
plan_base, plan_base,
download_into=temp_localname, download_into=temp_localname,
header_range_min=temp_localsize, header_range_min=temp_localsize,
@ -251,7 +253,7 @@ def prepare_plan(
plan_type='resume', plan_type='resume',
seek_to=temp_localsize, seek_to=temp_localsize,
) )
plan_partial = dict( plan_partial = dotdict.DotDict(
plan_base, plan_base,
download_into=real_localname, download_into=real_localname,
header_range_min=user_range_min, header_range_min=user_range_min,