summaryrefslogtreecommitdiffstats
path: root/eos_downloader/download.py
diff options
context:
space:
mode:
Diffstat (limited to 'eos_downloader/download.py')
-rw-r--r--eos_downloader/download.py77
1 files changed, 77 insertions, 0 deletions
diff --git a/eos_downloader/download.py b/eos_downloader/download.py
new file mode 100644
index 0000000..2297b04
--- /dev/null
+++ b/eos_downloader/download.py
@@ -0,0 +1,77 @@
+# flake8: noqa: F811
+# pylint: disable=unused-argument
+# pylint: disable=too-few-public-methods
+
+"""download module"""
+
+import os.path
+import signal
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+from typing import Iterable, Any
+
+import requests
+import rich
+from rich import console
+from rich.progress import (BarColumn, DownloadColumn, Progress, TaskID,
+ TextColumn, TimeElapsedColumn, TransferSpeedColumn)
+
+console = rich.get_console()
+done_event = Event()
+
+
+def handle_sigint(signum: Any, frame: Any) -> None:
+ """Progress bar handler"""
+ done_event.set()
+
+
+signal.signal(signal.SIGINT, handle_sigint)
+
+
+class DownloadProgressBar():
+ """
+ Object to manage Download process with Progress Bar from Rich
+ """
+
+ def __init__(self) -> None:
+ """
+ Class Constructor
+ """
+ self.progress = Progress(
+ TextColumn("💾 Downloading [bold blue]{task.fields[filename]}", justify="right"),
+ BarColumn(bar_width=None),
+ "[progress.percentage]{task.percentage:>3.1f}%",
+ "•",
+ TransferSpeedColumn(),
+ "•",
+ DownloadColumn(),
+ "•",
+ TimeElapsedColumn(),
+ "•",
+ console=console
+ )
+
+ def _copy_url(self, task_id: TaskID, url: str, path: str, block_size: int = 1024) -> bool:
+ """Copy data from a url to a local file."""
+ response = requests.get(url, stream=True, timeout=5)
+ # This will break if the response doesn't contain content length
+ self.progress.update(task_id, total=int(response.headers['Content-Length']))
+ with open(path, "wb") as dest_file:
+ self.progress.start_task(task_id)
+ for data in response.iter_content(chunk_size=block_size):
+ dest_file.write(data)
+ self.progress.update(task_id, advance=len(data))
+ if done_event.is_set():
+ return True
+ # console.print(f"Downloaded {path}")
+ return False
+
+ def download(self, urls: Iterable[str], dest_dir: str) -> None:
+ """Download multuple files to the given directory."""
+ with self.progress:
+ with ThreadPoolExecutor(max_workers=4) as pool:
+ for url in urls:
+ filename = url.split("/")[-1].split('?')[0]
+ dest_path = os.path.join(dest_dir, filename)
+ task_id = self.progress.add_task("download", filename=filename, start=False)
+ pool.submit(self._copy_url, task_id, url, dest_path)