Skip to content
Navigation Menu
{{ message }}
-
-
Notifications
You must be signed in to change notification settings - Fork 104
Expand file tree
/
Copy pathmulti.py
More file actions
410 lines (342 loc) · 13.3 KB
/
Copy pathmulti.py
File metadata and controls
410 lines (342 loc) · 13.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
from __future__ import annotations
import enum
import io
import itertools
import operator
import sys
import threading
import time
import timeit
import types
import typing
from datetime import timedelta
import python_utils
from . import bar, terminal
from .terminal import stream
SortKeyFunc = typing.Callable[[bar.ProgressBar], typing.Any]
class _Update(typing.Protocol):
def __call__(self, force: bool = True, write: bool = True) -> str: ...
class SortKey(str, enum.Enum):
"""
Sort keys for the MultiBar.
This is a string enum, so you can use any
progressbar attribute or property as a sort key.
Note that the multibar defaults to lazily rendering only the changed
progressbars. This means that sorting by dynamic attributes such as
`value` might result in more rendering which can have a small performance
impact.
"""
CREATED = 'index'
LABEL = 'label'
VALUE = 'value'
PERCENTAGE = 'percentage'
class MultiBar(dict[str, bar.ProgressBar]):
fd: typing.TextIO
_buffer: io.StringIO
#: The format for the label to append/prepend to the progressbar
label_format: str
#: Automatically prepend the label to the progressbars
prepend_label: bool
#: Automatically append the label to the progressbars
append_label: bool
#: If `initial_format` is `None`, the progressbar rendering is used
# which will *start* the progressbar. That means the progressbar will
# have no knowledge of your data and will run as an infinite progressbar.
initial_format: str | None
#: If `finished_format` is `None`, the progressbar rendering is used.
finished_format: str | None
#: The multibar updates at a fixed interval regardless of the progressbar
# updates
update_interval: float
remove_finished: float | None
#: The kwargs passed to the progressbar constructor
progressbar_kwargs: dict[str, typing.Any]
#: The progressbar sorting key function
sort_keyfunc: SortKeyFunc
_previous_output: list[str]
_finished_at: dict[bar.ProgressBar, float]
_labeled: set[bar.ProgressBar]
_print_lock: threading.RLock
_thread: threading.Thread | None
_thread_finished: threading.Event
_thread_closed: threading.Event
def __init__(
self,
bars: typing.Iterable[tuple[str, bar.ProgressBar]] | None = None,
fd: typing.TextIO = sys.stderr,
prepend_label: bool = True,
append_label: bool = False,
label_format: str = '{label:20.20} ',
initial_format: str | None = '{label:20.20} Not yet started',
finished_format: str | None = None,
update_interval: float = 1 / 60.0, # 60fps
show_initial: bool = True,
show_finished: bool = True,
remove_finished: timedelta | float = timedelta(seconds=3600),
sort_key: str | SortKey = SortKey.CREATED,
sort_reverse: bool = True,
sort_keyfunc: SortKeyFunc | None = None,
**progressbar_kwargs: typing.Any,
):
self.fd = fd
self.prepend_label = prepend_label
self.append_label = append_label
self.label_format = label_format
self.initial_format = initial_format
self.finished_format = finished_format
self.update_interval = update_interval
self.show_initial = show_initial
self.show_finished = show_finished
self.remove_finished = python_utils.delta_to_seconds_or_none(
remove_finished,
)
self.progressbar_kwargs = progressbar_kwargs
if sort_keyfunc is None:
sort_keyfunc = operator.attrgetter(sort_key)
self.sort_keyfunc = sort_keyfunc
self.sort_reverse = sort_reverse
self._labeled = set()
self._finished_at = {}
self._previous_output = []
self._buffer = io.StringIO()
self._print_lock = threading.RLock()
self._thread = None
self._thread_finished = threading.Event()
self._thread_closed = threading.Event()
super().__init__(bars or {})
def __setitem__(self, key: str, bar: bar.ProgressBar):
"""Add a progressbar to the multibar."""
if bar.label != key or not key: # pragma: no branch
bar.label = key
bar.fd = stream.LastLineStream(self.fd)
bar.paused = True
# Essentially `bar.print = self.print`, but `mypy` doesn't
# like that
bar.print = self.print # type: ignore
# Just in case someone is using a progressbar with a custom
# constructor and forgot to call the super constructor
if bar.index == -1:
bar.index = next(
bar._index_counter # pyright: ignore[reportPrivateUsage]
)
super().__setitem__(key, bar)
def __delitem__(self, key: str) -> None:
"""Remove a progressbar from the multibar."""
bar_: bar.ProgressBar = self.pop(key)
self._finished_at.pop(bar_, None)
self._labeled.discard(bar_)
def __getitem__(self, key: str):
"""Get (and create if needed) a progressbar from the multibar."""
try:
return super().__getitem__(key)
except KeyError:
progress = bar.ProgressBar(**self.progressbar_kwargs)
self[key] = progress
return progress
def _label_bar(self, bar: bar.ProgressBar) -> None:
if bar in self._labeled: # pragma: no branch
return
assert bar.widgets, 'Cannot prepend label to empty progressbar'
if self.prepend_label: # pragma: no branch
self._labeled.add(bar)
bar.widgets.insert(0, self.label_format.format(label=bar.label))
if self.append_label: # pragma: no branch
self._labeled.add(bar)
bar.widgets.append(self.label_format.format(label=bar.label))
def render(self, flush: bool = True, force: bool = False) -> None:
"""Render the multibar to the given stream."""
now: float = timeit.default_timer()
expired: float | None = (
now - self.remove_finished if self.remove_finished else None
)
# sourcery skip: list-comprehension
output: list[str] = []
for bar_ in self.get_sorted_bars():
if not bar_.started() and not self.show_initial:
continue
output.extend(
iter(self._render_bar(bar_, expired=expired, now=now)),
)
with self._print_lock:
# Clear the previous output if progressbars have been removed
for i in range(len(output), len(self._previous_output)):
self._buffer.write(
terminal.clear_line(i + 1),
) # pragma: no cover
# Add empty lines to the end of the output if progressbars have
# been added
for _ in range(len(self._previous_output), len(output)):
# Adding a new line so we don't overwrite previous output
self._buffer.write('\n')
for i, (previous, current) in enumerate(
itertools.zip_longest(
self._previous_output,
output,
fillvalue='',
),
):
if previous != current or force: # pragma: no branch
self.print(
'\r' + current.strip(),
offset=i + 1,
end='',
clear=False,
flush=False,
)
self._previous_output = output
if flush: # pragma: no branch
self.flush()
def _render_bar(
self,
bar_: bar.ProgressBar,
now: float,
expired: float | None,
) -> typing.Iterable[str]:
def update(
force: bool = True, write: bool = True
) -> str: # pragma: no cover
self._label_bar(bar_)
bar_.update(force=force)
if write:
return typing.cast(stream.LastLineStream, bar_.fd).line
else:
return ''
if bar_.finished():
yield from self._render_finished_bar(bar_, now, expired, update)
elif bar_.started():
update()
else:
if self.initial_format is None:
bar_.start()
yield update()
else:
yield self.initial_format.format(label=bar_.label)
def _render_finished_bar(
self,
bar_: bar.ProgressBar,
now: float,
expired: float | None,
update: _Update,
) -> typing.Iterable[str]:
if bar_ not in self._finished_at:
self._finished_at[bar_] = now
# Force update to get the finished format
update(write=False)
if (
self.remove_finished
and expired is not None
and expired >= self._finished_at[bar_]
):
del self[bar_.label]
return
if not self.show_finished:
return
if bar_.finished(): # pragma: no branch
if self.finished_format is None:
update(force=False)
else: # pragma: no cover
yield self.finished_format.format(label=bar_.label)
def print(
self,
*args: typing.Any,
end: str = '\n',
offset: int | None = None,
flush: bool = True,
clear: bool = True,
**kwargs: typing.Any,
):
"""
Print to the progressbar stream without overwriting the progressbars.
Args:
end: The string to append to the end of the output
offset: The number of lines to offset the output by. If None, the
output will be printed above the progressbars
flush: Whether to flush the output to the stream
clear: If True, the line will be cleared before printing.
**kwargs: Additional keyword arguments to pass to print
"""
with self._print_lock:
if offset is None:
offset = len(self._previous_output)
if not clear:
self._buffer.write(terminal.PREVIOUS_LINE(offset))
if clear:
self._buffer.write(terminal.PREVIOUS_LINE(offset))
self._buffer.write(terminal.CLEAR_LINE_ALL())
print(*args, **kwargs, file=self._buffer, end=end)
if clear:
self._buffer.write(terminal.CLEAR_SCREEN_TILL_END())
for line in self._previous_output:
self._buffer.write(line.strip())
self._buffer.write('\n')
else:
self._buffer.write(terminal.NEXT_LINE(offset))
if flush:
self.flush()
def flush(self) -> None:
# The fd write happens under the lock as well so concurrent
# print()/render() calls cannot interleave their output
with self._print_lock:
value = self._buffer.getvalue()
self._buffer.seek(0)
self._buffer.truncate(0)
self.fd.write(value)
self.fd.flush()
def run(self, join: bool = True) -> None:
"""
Start the multibar render loop and run the progressbars until they
have force _thread_finished.
"""
while not self._thread_finished.is_set(): # pragma: no branch
self.render()
time.sleep(self.update_interval)
if join or self._thread_closed.is_set():
# If the thread is closed, we need to check if the progressbars
# have finished. If they have, we can exit the loop
for bar_ in list(self.values()): # pragma: no cover
if not bar_.finished():
break
else:
# Render one last time to make sure the progressbars are
# correctly finished
self.render(force=True)
return
def start(self) -> None:
assert not self._thread, 'Multibar already started'
self._thread_finished.clear()
self._thread_closed.clear()
self._thread = threading.Thread(
target=self.run,
args=(False,),
daemon=True,
)
self._thread.start()
def join(self, timeout: float | None = None) -> None:
if self._thread is not None:
self._thread_closed.set()
self._thread.join(timeout=timeout)
if not self._thread.is_alive():
self._thread = None
def stop(self, timeout: float | None = None):
self._thread_finished.set()
self.join(timeout=timeout)
def get_sorted_bars(self):
# Materialize the values into a list first so other threads can
# add or remove bars while we are sorting and rendering
bars = list(self.values())
return sorted(bars, key=self.sort_keyfunc, reverse=self.sort_reverse)
def __enter__(self):
self.start()
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: types.TracebackType | None,
) -> bool | None:
if exc_type is None:
self.join()
else:
# Don't wait for unfinished progressbars when an exception is
# propagating; that would block forever
self.stop()
You can’t perform that action at this time.