]>
Commit | Line | Data |
---|---|---|
1c51c520 SS |
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 |