Some threaded_dl improvements.

This commit is contained in:
voussoir 2022-11-02 18:42:14 -07:00
parent d92f190590
commit 8124477398
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB

View file

@ -1,6 +1,7 @@
import argparse import argparse
import ast import ast
import os import os
import random
import shutil import shutil
import sys import sys
import threading import threading
@ -78,15 +79,16 @@ def normalize_headers(headers):
return {key: val for (key, val) in zip(keys, vals)} return {key: val for (key, val) in zip(keys, vals)}
if isinstance(headers, str) and os.path.isfile(headers): if isinstance(headers, str) and os.path.isfile(headers):
headers = pathclass.Path(headers).readlines('r', encoding='utf-8') headers = pathclass.Path(headers).read('r', encoding='utf-8')
if isinstance(headers, str): if isinstance(headers, str):
if headers.startswith('{'): if headers.startswith('{'):
return ast.literal_eval(headers) return ast.literal_eval(headers)
else: else:
lines = [line for line in headers.splitlines() if line.strip()] lines = [line for line in headers.splitlines() if line.strip()]
pairs = [line.strip().split(': ', 1) for line in lines] lines = [line for line in lines if not line.startswith('#')]
return {key: value for (key, value) in pairs} pairs = [line.strip().split(':', 1) for line in lines]
return {key.strip(): value.strip() for (key, value) in pairs}
return headers return headers
@ -109,6 +111,7 @@ def prepare_urls_filenames(urls, filename_format):
if isinstance(url, (tuple, list)): if isinstance(url, (tuple, list)):
(url, filename) = url (url, filename) = url
else: else:
index1 = index + 1
basename = downloady.basename_from_url(url) basename = downloady.basename_from_url(url)
extension = os.path.splitext(basename)[1] extension = os.path.splitext(basename)[1]
filename = filename_format.format( filename = filename_format.format(
@ -116,6 +119,7 @@ def prepare_urls_filenames(urls, filename_format):
ext=extension, ext=extension,
extension=extension, extension=extension,
index=index, index=index,
index1=index1,
now=now, now=now,
) )
@ -179,7 +183,7 @@ def threaded_dl(
status = 0 status = 0
for job in pool.result_generator(): for job in pool.result_generator():
if job.exception: if job.exception:
log.error(traceback.format_exc()) log.error(''.join(traceback.format_exception(None, job.exception, job.exception.__traceback__)))
status = 1 status = 1
ui_stop_event.set() ui_stop_event.set()
@ -204,6 +208,7 @@ def threaded_dl_argparse(args):
urls = [u.split(' ', 1) if ' ' in u else u for u in urls] urls = [u.split(' ', 1) if ' ' in u else u for u in urls]
headers = normalize_headers(args.headers) headers = normalize_headers(args.headers)
print(headers)
bytespersecond = args.bytespersecond bytespersecond = args.bytespersecond
if bytespersecond is not None: if bytespersecond is not None:
@ -272,6 +277,19 @@ def main(argv):
) )
parser.add_argument( parser.add_argument(
'--headers', nargs='+', default=None, '--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) parser.set_defaults(func=threaded_dl_argparse)