Files
ezzdl/download.py
BarsTiger 2203874ca5 1.0.0
2022-05-06 11:54:59 +03:00

71 lines
2.0 KiB
Python

# Idea from rich examples
# https://github.com/Textualize/rich/blob/master/examples/downloader.py
import os.path
from concurrent.futures import ThreadPoolExecutor
import signal
from functools import partial
from threading import Event
from urllib.request import urlopen
from urllib.parse import unquote
from time import strftime
from rich.progress import (
BarColumn,
DownloadColumn,
Progress,
TaskID,
TextColumn,
TimeRemainingColumn,
TransferSpeedColumn,
)
from rich import print
progress = Progress(
TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
BarColumn(bar_width=None),
"[progress.percentage]{task.percentage:>3.1f}%",
"",
DownloadColumn(),
"",
TransferSpeedColumn(),
"",
TimeRemainingColumn(),
)
done_event = Event()
def handle_sigint(signum, frame):
done_event.set()
signal.signal(signal.SIGINT, handle_sigint)
def copy_url(task_id: TaskID, url: str, path: str) -> None:
print(f"[#1d4b6e][{strftime('%H:%M:%S')}][/] Requesting {url}")
response = urlopen(url)
progress.update(task_id, total=int(response.info()["Content-length"]))
with open(path, "wb") as dest_file:
progress.start_task(task_id)
for data in iter(partial(response.read, 32768), b""):
dest_file.write(data)
progress.update(task_id, advance=len(data))
if done_event.is_set():
return
print(f"[#1d4b6e][{strftime('%H:%M:%S')}][/] [#8dc789]Downloaded {os.path.abspath(path)}[/]")
def dl(urls, dest_dir: str):
if not os.path.isdir(dest_dir):
os.makedirs(dest_dir, exist_ok=True)
with progress:
with ThreadPoolExecutor(max_workers=len(urls)) as pool:
for url in urls:
filename = unquote(url.split("/")[-1])
dest_path = os.path.join(dest_dir, filename)
task_id = progress.add_task("download", filename=filename, start=False)
pool.submit(copy_url, task_id, url, dest_path)