From 85631c8ce514b487c9fb706b38cdfacc1a846adf Mon Sep 17 00:00:00 2001 From: Ivan Habunek Date: Sat, 13 Aug 2022 09:38:54 +0200 Subject: [PATCH] Extract progress tracking --- tests/test_progress.py | 94 ++++++++++++++++++++++++++++++++++++ twitchdl/progress.py | 106 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+) create mode 100644 tests/test_progress.py create mode 100644 twitchdl/progress.py diff --git a/tests/test_progress.py b/tests/test_progress.py new file mode 100644 index 0000000..dcb711f --- /dev/null +++ b/tests/test_progress.py @@ -0,0 +1,94 @@ +from twitchdl.progress import Progress + + +def test_initial_values(): + progress = Progress(10) + assert progress.downloaded == 0 + assert progress.estimated_total is None + assert progress.progress_perc == 0 + assert progress.remaining_time is None + assert progress.speed is None + assert progress.vod_count == 10 + assert progress.vod_downloaded_count == 0 + + +def test_downloaded(): + progress = Progress(3) + progress.start(1, 300) + progress.start(2, 300) + progress.start(3, 300) + + assert progress.downloaded == 0 + assert progress.progress_perc == 0 + + progress.advance(1, 100) + assert progress.downloaded == 100 + assert progress.progress_perc == 11 + + progress.advance(2, 200) + assert progress.downloaded == 300 + assert progress.progress_perc == 33 + + progress.advance(3, 150) + assert progress.downloaded == 450 + assert progress.progress_perc == 50 + + progress.advance(1, 50) + assert progress.downloaded == 500 + assert progress.progress_perc == 55 + + progress.abort(2) + assert progress.downloaded == 300 + assert progress.progress_perc == 33 + + progress.start(2, 300) + + progress.advance(1, 150) + progress.advance(2, 300) + progress.advance(3, 150) + + assert progress.downloaded == 900 + assert progress.progress_perc == 100 + + progress.end(1) + progress.end(2) + progress.end(3) + + assert progress.downloaded == 900 + assert progress.progress_perc == 100 + + +def test_estimated_total(): + progress = Progress(3) + assert progress.estimated_total is None + + progress.start(1, 12000) + assert progress.estimated_total == 12000 * 3 + + progress.start(2, 11000) + assert progress.estimated_total == 11500 * 3 + + progress.start(3, 10000) + assert progress.estimated_total == 11000 * 3 + + +def test_vod_downloaded_count(): + progress = Progress(3) + + progress.start(1, 100) + progress.start(2, 100) + progress.start(3, 100) + + assert progress.vod_downloaded_count == 0 + + progress.advance(1, 100) + progress.end(1) + assert progress.vod_downloaded_count == 1 + + progress.advance(2, 100) + progress.end(2) + assert progress.vod_downloaded_count == 2 + + progress.advance(3, 100) + progress.end(3) + assert progress.vod_downloaded_count == 3 diff --git a/twitchdl/progress.py b/twitchdl/progress.py new file mode 100644 index 0000000..5ce57c6 --- /dev/null +++ b/twitchdl/progress.py @@ -0,0 +1,106 @@ +import logging +import time + +from dataclasses import dataclass, field +from statistics import mean +from typing import Dict, Optional + +from twitchdl.output import print_out +from twitchdl.utils import format_size, format_duration + +logger = logging.getLogger(__name__) + + +TaskId = int + + +@dataclass +class Task: + id: TaskId + size: int + downloaded: int = 0 + + def advance(self, size): + self.downloaded += size + + +@dataclass +class Progress: + vod_count: int + downloaded: int = 0 + estimated_total: Optional[int] = None + progress_perc: int = 0 + remaining_time: Optional[int] = None + speed: Optional[float] = None + start_time: float = field(default_factory=time.time) + tasks: Dict[TaskId, Task] = field(default_factory=dict) + vod_downloaded_count: int = 0 + + def start(self, task_id: int, size: int): + logger.debug(f"#{task_id} start {size}b") + + if task_id in self.tasks: + raise ValueError(f"Task {task_id}: cannot start, already started") + + self.tasks[task_id] = Task(task_id, size) + self._calculate_total() + self._calculate_progress() + self.print() + + def advance(self, task_id: int, chunk_size: int): + logger.debug(f"#{task_id} advance {chunk_size}") + + if task_id not in self.tasks: + raise ValueError(f"Task {task_id}: cannot advance, not started") + + self.downloaded += chunk_size + self.tasks[task_id].advance(chunk_size) + self._calculate_progress() + self.print() + + def abort(self, task_id: int): + logger.debug(f"#{task_id} abort") + + if task_id not in self.tasks: + raise ValueError(f"Task {task_id}: cannot abort, not started") + + del self.tasks[task_id] + self.downloaded = sum(t.downloaded for t in self.tasks.values()) + + self._calculate_total() + self._calculate_progress() + self.print() + + def end(self, task_id: int): + logger.debug(f"#{task_id} end") + + if task_id not in self.tasks: + raise ValueError(f"Task {task_id}: cannot end, not started") + + task = self.tasks[task_id] + if task.size != task.downloaded: + logger.warn(f"Taks {task_id} ended with {task.downloaded}b downloaded, expected {task.size}b.") + + self.vod_downloaded_count += 1 + self.print() + + def _calculate_total(self): + self.estimated_total = int(mean(t.size for t in self.tasks.values()) * self.vod_count) if self.tasks else None + + def _calculate_progress(self): + elapsed_time = time.time() - self.start_time + self.progress_perc = int(100 * self.downloaded / self.estimated_total) if self.estimated_total else 0 + self.speed = self.downloaded / elapsed_time if elapsed_time else None + self.remaining_time = int((self.estimated_total - self.downloaded) / self.speed) if self.estimated_total and self.speed else None + + def print(self): + progress = " ".join([ + f"Downloaded {self.vod_downloaded_count}/{self.vod_count} VODs", + f"({self.progress_perc}%)", + f"{format_size(self.downloaded)}", + f"of ~{format_size(self.estimated_total)}" if self.estimated_total else "", + f"at {format_size(self.speed)}/s" if self.speed else "", + f"remaining ~{format_duration(self.remaining_time)}" if self.remaining_time is not None else "", + ]) + + print_out(f"\r{progress} ", end="")