summaryrefslogtreecommitdiffstats
path: root/yt_dlp/utils/progress.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--yt_dlp/utils/progress.py109
1 files changed, 109 insertions, 0 deletions
diff --git a/yt_dlp/utils/progress.py b/yt_dlp/utils/progress.py
new file mode 100644
index 0000000..f254a38
--- /dev/null
+++ b/yt_dlp/utils/progress.py
@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import bisect
+import threading
+import time
+
+
+class ProgressCalculator:
+ # Time to calculate the speed over (seconds)
+ SAMPLING_WINDOW = 3
+ # Minimum timeframe before to sample next downloaded bytes (seconds)
+ SAMPLING_RATE = 0.05
+ # Time before showing eta (seconds)
+ GRACE_PERIOD = 1
+
+ def __init__(self, initial: int):
+ self._initial = initial or 0
+ self.downloaded = self._initial
+
+ self.elapsed: float = 0
+ self.speed = SmoothValue(0, smoothing=0.7)
+ self.eta = SmoothValue(None, smoothing=0.9)
+
+ self._total = 0
+ self._start_time = time.monotonic()
+ self._last_update = self._start_time
+
+ self._lock = threading.Lock()
+ self._thread_sizes: dict[int, int] = {}
+
+ self._times = [self._start_time]
+ self._downloaded = [self.downloaded]
+
+ @property
+ def total(self):
+ return self._total
+
+ @total.setter
+ def total(self, value: int | None):
+ with self._lock:
+ if value is not None and value < self.downloaded:
+ value = self.downloaded
+
+ self._total = value
+
+ def thread_reset(self):
+ current_thread = threading.get_ident()
+ with self._lock:
+ self._thread_sizes[current_thread] = 0
+
+ def update(self, size: int | None):
+ if not size:
+ return
+
+ current_thread = threading.get_ident()
+
+ with self._lock:
+ last_size = self._thread_sizes.get(current_thread, 0)
+ self._thread_sizes[current_thread] = size
+ self._update(size - last_size)
+
+ def _update(self, size: int):
+ current_time = time.monotonic()
+
+ self.downloaded += size
+ self.elapsed = current_time - self._start_time
+ if self.total is not None and self.downloaded > self.total:
+ self._total = self.downloaded
+
+ if self._last_update + self.SAMPLING_RATE > current_time:
+ return
+ self._last_update = current_time
+
+ self._times.append(current_time)
+ self._downloaded.append(self.downloaded)
+
+ offset = bisect.bisect_left(self._times, current_time - self.SAMPLING_WINDOW)
+ del self._times[:offset]
+ del self._downloaded[:offset]
+ if len(self._times) < 2:
+ self.speed.reset()
+ self.eta.reset()
+ return
+
+ download_time = current_time - self._times[0]
+ if not download_time:
+ return
+
+ self.speed.set((self.downloaded - self._downloaded[0]) / download_time)
+ if self.total and self.speed.value and self.elapsed > self.GRACE_PERIOD:
+ self.eta.set((self.total - self.downloaded) / self.speed.value)
+ else:
+ self.eta.reset()
+
+
+class SmoothValue:
+ def __init__(self, initial: float | None, smoothing: float):
+ self.value = self.smooth = self._initial = initial
+ self._smoothing = smoothing
+
+ def set(self, value: float):
+ self.value = value
+ if self.smooth is None:
+ self.smooth = self.value
+ else:
+ self.smooth = (1 - self._smoothing) * value + self._smoothing * self.smooth
+
+ def reset(self):
+ self.value = self.smooth = self._initial