yt_dlp/utils/progress.py

   1 from __future__ import annotations
   2
   3 import bisect
   4 import threading
   5 import time
   6
   7
   8 class ProgressCalculator:
   9     # Time to calculate the speed over (seconds)
  10     SAMPLING_WINDOW = 3
  11     # Minimum timeframe before to sample next downloaded bytes (seconds)
  12     SAMPLING_RATE = 0.05
  13     # Time before showing eta (seconds)
  14     GRACE_PERIOD = 1
  15
  16     def __init__(self, initial: int):
  17         self._initial = initial or 0
  18         self.downloaded = self._initial
  19
  20         self.elapsed: float = 0
  21         self.speed = SmoothValue(0, smoothing=0.7)
  22         self.eta = SmoothValue(None, smoothing=0.9)
  23
  24         self._total = 0
  25         self._start_time = time.monotonic()
  26         self._last_update = self._start_time
  27
  28         self._lock = threading.Lock()
  29         self._thread_sizes: dict[int, int] = {}
  30
  31         self._times = [self._start_time]
  32         self._downloaded = [self.downloaded]
  33
  34     @property
  35     def total(self):
  36         return self._total
  37
  38     @total.setter
  39     def total(self, value: int | None):
  40         with self._lock:
  41             if value is not None and value < self.downloaded:
  42                 value = self.downloaded
  43
  44             self._total = value
  45
  46     def thread_reset(self):
  47         current_thread = threading.get_ident()
  48         with self._lock:
  49             self._thread_sizes[current_thread] = 0
  50
  51     def update(self, size: int | None):
  52         if not size:
  53             return
  54
  55         current_thread = threading.get_ident()
  56
  57         with self._lock:
  58             last_size = self._thread_sizes.get(current_thread, 0)
  59             self._thread_sizes[current_thread] = size
  60             self._update(size - last_size)
  61
  62     def _update(self, size: int):
  63         current_time = time.monotonic()
  64
  65         self.downloaded += size
  66         self.elapsed = current_time - self._start_time
  67         if self.total is not None and self.downloaded > self.total:
  68             self._total = self.downloaded
  69
  70         if self._last_update + self.SAMPLING_RATE > current_time:
  71             return
  72         self._last_update = current_time
  73
  74         self._times.append(current_time)
  75         self._downloaded.append(self.downloaded)
  76
  77         offset = bisect.bisect_left(self._times, current_time - self.SAMPLING_WINDOW)
  78         del self._times[:offset]
  79         del self._downloaded[:offset]
  80         if len(self._times) < 2:
  81             self.speed.reset()
  82             self.eta.reset()
  83             return
  84
  85         download_time = current_time - self._times[0]
  86         if not download_time:
  87             return
  88
  89         self.speed.set((self.downloaded - self._downloaded[0]) / download_time)
  90         if self.total and self.speed.value and self.elapsed > self.GRACE_PERIOD:
  91             self.eta.set((self.total - self.downloaded) / self.speed.value)
  92         else:
  93             self.eta.reset()
  94
  95
  96 class SmoothValue:
  97     def __init__(self, initial: float | None, smoothing: float):
  98         self.value = self.smooth = self._initial = initial
  99         self._smoothing = smoothing
 100
 101     def set(self, value: float):
 102         self.value = value
 103         if self.smooth is None:
 104             self.smooth = self.value
 105         else:
 106             self.smooth = (1 - self._smoothing) * value + self._smoothing * self.smooth
 107
 108     def reset(self):
 109         self.value = self.smooth = self._initial