301 lines
8.9 KiB
Python
301 lines
8.9 KiB
Python
import argparse
|
|
import ast
|
|
import os
|
|
import random
|
|
import shutil
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
|
|
from voussoirkit import betterhelp
|
|
from voussoirkit import bytestring
|
|
from voussoirkit import downloady
|
|
from voussoirkit import pathclass
|
|
from voussoirkit import pipeable
|
|
from voussoirkit import ratelimiter
|
|
from voussoirkit import ratemeter
|
|
from voussoirkit import threadpool
|
|
from voussoirkit import vlogging
|
|
|
|
log = vlogging.getLogger(__name__, 'threaded_dl')
|
|
downloady.log.setLevel(vlogging.WARNING)
|
|
|
|
def clean_url_list(urls):
|
|
for url in urls:
|
|
if isinstance(url, (tuple, list)):
|
|
(url, filename) = url
|
|
else:
|
|
filename = None
|
|
url = url.strip()
|
|
|
|
if not url:
|
|
continue
|
|
|
|
if url.startswith('#'):
|
|
continue
|
|
|
|
if filename:
|
|
yield (url, filename)
|
|
else:
|
|
yield url
|
|
|
|
def download_job(
|
|
url,
|
|
filename,
|
|
*,
|
|
bytespersecond=None,
|
|
headers=None,
|
|
meter=None,
|
|
timeout=None,
|
|
):
|
|
log.info(f'Starting "{filename}"')
|
|
downloady.download_file(
|
|
url,
|
|
filename,
|
|
bytespersecond=bytespersecond,
|
|
headers=headers,
|
|
ratemeter=meter,
|
|
timeout=timeout,
|
|
)
|
|
log.info(f'Finished "{filename}"')
|
|
|
|
def normalize_headers(headers):
|
|
if headers is None:
|
|
return {}
|
|
|
|
if not headers:
|
|
return {}
|
|
|
|
if isinstance(headers, dict):
|
|
return headers
|
|
|
|
if isinstance(headers, list) and len(headers) == 1:
|
|
headers = headers[0]
|
|
|
|
if isinstance(headers, (list, tuple)):
|
|
keys = headers[::2]
|
|
vals = headers[1::2]
|
|
return {key: val for (key, val) in zip(keys, vals)}
|
|
|
|
if isinstance(headers, str) and os.path.isfile(headers):
|
|
headers = pathclass.Path(headers).read('r', encoding='utf-8')
|
|
|
|
if isinstance(headers, str):
|
|
if headers.startswith('{'):
|
|
return ast.literal_eval(headers)
|
|
else:
|
|
lines = [line for line in headers.splitlines() if line.strip()]
|
|
lines = [line for line in lines if not line.startswith('#')]
|
|
pairs = [line.strip().split(':', 1) for line in lines]
|
|
return {key.strip(): value.strip() for (key, value) in pairs}
|
|
|
|
return headers
|
|
|
|
def prepare_urls_filenames(urls, filename_format):
|
|
now = int(time.time())
|
|
|
|
if os.path.normcase(filename_format) != os.devnull:
|
|
index_digits = len(str(len(urls)))
|
|
filename_format = filename_format.replace('{index}', '{index:0%0dd}' % index_digits)
|
|
|
|
if '{' not in filename_format and len(urls) > 1:
|
|
filename_format += '_{index}'
|
|
|
|
if '{extension}' not in filename_format and '{basename}' not in filename_format:
|
|
filename_format += '{extension}'
|
|
|
|
urls_filenames = []
|
|
|
|
for (index, url) in enumerate(clean_url_list(urls)):
|
|
if isinstance(url, (tuple, list)):
|
|
(url, filename) = url
|
|
else:
|
|
index1 = index + 1
|
|
basename = downloady.basename_from_url(url)
|
|
extension = os.path.splitext(basename)[1]
|
|
filename = filename_format.format(
|
|
basename=basename,
|
|
ext=extension,
|
|
extension=extension,
|
|
index=index,
|
|
index1=index1,
|
|
now=now,
|
|
random=random.getrandbits(32),
|
|
)
|
|
|
|
if os.path.exists(filename):
|
|
log.info(f'Skipping existing file "{filename}"')
|
|
continue
|
|
|
|
urls_filenames.append((url, filename))
|
|
|
|
return urls_filenames
|
|
|
|
def threaded_dl(
|
|
urls,
|
|
thread_count,
|
|
filename_format,
|
|
bytespersecond=None,
|
|
headers=None,
|
|
timeout=None,
|
|
):
|
|
urls_filenames = prepare_urls_filenames(urls, filename_format)
|
|
|
|
if not urls_filenames:
|
|
return
|
|
|
|
if bytespersecond is not None:
|
|
# It is important that we convert this to a Ratelimter now instead of
|
|
# passing the user's integer to downloady, because we want all threads
|
|
# to share a single limiter instance instead of each creating their
|
|
# own by the integer.
|
|
bytespersecond = ratelimiter.Ratelimiter(bytespersecond)
|
|
|
|
meter = ratemeter.RateMeter(span=5)
|
|
|
|
pool = threadpool.ThreadPool(thread_count, paused=True)
|
|
|
|
ui_stop_event = threading.Event()
|
|
ui_kwargs = {
|
|
'meter': meter,
|
|
'stop_event': ui_stop_event,
|
|
'pool': pool,
|
|
}
|
|
ui_thread = threading.Thread(target=ui_thread_func, kwargs=ui_kwargs, daemon=True)
|
|
ui_thread.start()
|
|
|
|
kwargss = []
|
|
for (url, filename) in urls_filenames:
|
|
kwargs = {
|
|
'function': download_job,
|
|
'kwargs': {
|
|
'bytespersecond': bytespersecond,
|
|
'filename': filename,
|
|
'headers': headers,
|
|
'meter': meter,
|
|
'timeout': timeout,
|
|
'url': url,
|
|
}
|
|
}
|
|
kwargss.append(kwargs)
|
|
pool.add_many(kwargss)
|
|
|
|
status = 0
|
|
for job in pool.result_generator():
|
|
if job.exception:
|
|
log.error(''.join(traceback.format_exception(None, job.exception, job.exception.__traceback__)))
|
|
status = 1
|
|
|
|
ui_stop_event.set()
|
|
ui_thread.join()
|
|
return status
|
|
|
|
def ui_thread_func(meter, pool, stop_event):
|
|
if pipeable.stdout_pipe():
|
|
return
|
|
|
|
while not stop_event.is_set():
|
|
width = shutil.get_terminal_size().columns
|
|
speed = meter.report()[2]
|
|
message = f'{bytestring.bytestring(speed)}/s | {pool.running_count} threads'
|
|
spaces = ' ' * (width - len(message) - 1)
|
|
pipeable.stderr(message + spaces, end='\r')
|
|
|
|
stop_event.wait(timeout=0.5)
|
|
|
|
def threaded_dl_argparse(args):
|
|
urls = pipeable.input(args.url_file, read_files=True, skip_blank=True, strip=True)
|
|
urls = [u.split(' ', 1) if ' ' in u else u for u in urls]
|
|
|
|
headers = normalize_headers(args.headers)
|
|
print(headers)
|
|
|
|
bytespersecond = args.bytespersecond
|
|
if bytespersecond is not None:
|
|
bytespersecond = bytestring.parsebytes(bytespersecond)
|
|
|
|
return threaded_dl(
|
|
urls,
|
|
bytespersecond=bytespersecond,
|
|
filename_format=args.filename_format,
|
|
headers=headers,
|
|
thread_count=args.thread_count,
|
|
timeout=args.timeout,
|
|
)
|
|
|
|
@vlogging.main_decorator
|
|
def main(argv):
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
'url_file',
|
|
metavar='links',
|
|
help='''
|
|
The name of a file containing links to download, one per line.
|
|
Uses pipeable to support !c clipboard, !i stdin lines of urls.
|
|
''',
|
|
)
|
|
parser.add_argument(
|
|
'thread_count',
|
|
type=int,
|
|
help='''
|
|
Integer number of threads to use for downloading.
|
|
''',
|
|
)
|
|
parser.add_argument(
|
|
'filename_format',
|
|
nargs='?',
|
|
type=str,
|
|
default='{now}_{index}_{basename}',
|
|
help='''
|
|
A string that controls the names of the downloaded files. Uses Python's
|
|
brace-style formatting. Available formatters are:
|
|
- {basename}: The name of the file as indicated by the URL.
|
|
E.g. example.com/image.jpg -> image.jpg
|
|
- {extension}: The extension of the file as indicated by the URL, including
|
|
the dot. E.g. example.com/image.jpg -> .jpg
|
|
- {index}: The index of this URL within the sequence of all downloaded URLs.
|
|
Starts from 0.
|
|
- {now}: The unix timestamp at which this download job was started. It might
|
|
be ugly but at least it's unambiguous when doing multiple download batches
|
|
with similar filenames.
|
|
''',
|
|
)
|
|
parser.add_argument(
|
|
'--bytespersecond',
|
|
default=None,
|
|
help='''
|
|
Limit the overall download speed to X bytes per second. Uses
|
|
bytestring.parsebytes to support strings like "1m", "500k", "2 mb", etc.
|
|
''',
|
|
)
|
|
parser.add_argument(
|
|
'--timeout',
|
|
default=15,
|
|
help='''
|
|
Integer number of seconds to use as HTTP request timeout for each download.
|
|
''',
|
|
)
|
|
parser.add_argument(
|
|
'--headers', nargs='+', default=None,
|
|
help='''
|
|
HTTP headers to add to your request. There are many ways to specify headers:
|
|
|
|
You can provide multiple command line arguments where the first is a key,
|
|
the second is its value, the third is another key, the fourth is its value...
|
|
|
|
You can provide a single command line argument which is a JSON string containing
|
|
key:value pairs.
|
|
|
|
You can provide a single command line argument which is a filename.
|
|
The file can be a JSON file, or alternatively the file should have each
|
|
key:value on a separate line and a colon should separate each key from its value.
|
|
''',
|
|
)
|
|
parser.set_defaults(func=threaded_dl_argparse)
|
|
|
|
return betterhelp.go(parser, argv)
|
|
|
|
if __name__ == '__main__':
|
|
raise SystemExit(main(sys.argv[1:]))
|