rich/examples/downloader.py
2020-03-08 16:49:36 +00:00

40 lines
1.4 KiB
Python

from concurrent.futures import ThreadPoolExecutor
from functools import partial
import os.path
import sys
from typing import Iterable
from urllib.request import urlopen
from rich.progress import bar_widget, Progress, TaskID
progress = Progress(bar_widget, "{task.percentage:0.1f}%", "{task.fields[filename]}")
def copy_url(task_id: TaskID, url: str, path: str) -> None:
"""Copy data from a url to a local file."""
response = urlopen(url)
# This will break if the response doesn't content content length
progress.update(task_id, total=int(response.info()["Content-length"]))
with open(path, "wb") as dest_file:
for data in iter(partial(response.read, 32768), b""):
dest_file.write(data)
progress.update(task_id, advance=len(data))
def download(urls: Iterable[str], dest_dir: str):
"""Download multuple files to the given directory."""
with progress:
with ThreadPoolExecutor(max_workers=4) as pool:
for url in urls:
filename = url.split("/")[-1]
task_id = progress.add_task("download", filename=filename)
dest_path = os.path.join(dest_dir, filename)
pool.submit(copy_url, task_id, url, dest_path)
if __name__ == "__main__":
if sys.argv[1:]:
download(sys.argv[1:], "./")
else:
print("Usage:\npython downloader.py URL1 URL2 URL3 (etc)")