Extract progress tracking

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

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"<cyan>{format_size(self.downloaded)}</cyan>",
+            f"of <cyan>~{format_size(self.estimated_total)}</cyan>" if self.estimated_total else "",
+            f"at <cyan>{format_size(self.speed)}/s</cyan>" if self.speed else "",
+            f"remaining <cyan>~{format_duration(self.remaining_time)}</cyan>" if self.remaining_time is not None else "",
+        ])
+
+        print_out(f"\r{progress}     ", end="")