1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import sys
7import time
8import logging
10from .get_terminal_size import get_terminal_size
12logger = logging.getLogger('pyrocko.progress')
14# TODO: Refactor so that if multiple viewers are attached, they can do
15# their updates independently of each other (at independent time intervals).
17# TODO: Refactor so that different viewers can render task states differently.
19# spinner = u'\u25dc\u25dd\u25de\u25df'
20# spinner = '⣾⣽⣻⢿⡿⣟⣯⣷'
21spinner = "◴◷◶◵"
22skull = u'\u2620'
23check = u'\u2714'
24cross = u'\u2716'
25bar = u'[- ]'
26blocks = u'\u2588\u2589\u258a\u258b\u258c\u258d\u258e\u258f '
28symbol_done = check
29symbol_failed = cross # skull
31ansi_up = u'\033[%iA'
32ansi_down = u'\033[%iB'
33ansi_left = u'\033[%iC'
34ansi_right = u'\033[%iD'
35ansi_next_line = u'\033E'
37ansi_erase_display = u'\033[2J'
38ansi_window = u'\033[%i;%ir'
39ansi_move_to = u'\033[%i;%iH'
41ansi_clear_down = u'\033[0J'
42ansi_clear_up = u'\033[1J'
43ansi_clear = u'\033[2J'
45ansi_clear_right = u'\033[0K'
47ansi_scroll_up = u'\033D'
48ansi_scroll_down = u'\033M'
50ansi_reset = u'\033c'
53g_force_viewer_off = False
55g_viewer = 'terminal'
58def set_default_viewer(viewer):
59 global g_viewer
60 g_viewer = viewer
63class StatusViewer(object):
65 def __init__(self, parent=None):
66 self._parent = parent
68 def __enter__(self):
69 return self
71 def __exit__(self, *_):
72 self.stop()
74 def stop(self):
75 if self._parent:
76 self._parent.hide(self)
78 def draw(self, lines):
79 pass
82class TerminalStatusViewer(StatusViewer):
83 def __init__(self, parent=None):
84 self._terminal_size = get_terminal_size()
85 self._height = 0
86 self._state = 0
87 self._parent = parent
89 def print(self, s):
90 print(s, end='', file=sys.stderr)
92 def flush(self):
93 print('', end='', flush=True, file=sys.stderr)
95 def start(self):
96 sx, sy = self._terminal_size
97 self._state = 1
99 def stop(self):
100 if self._state == 1:
101 sx, sy = self._terminal_size
102 self._resize(0)
103 self.print(ansi_move_to % (sy-self._height, 1))
104 self.flush()
106 self._state = 2
107 if self._parent:
108 self._parent.hide(self)
110 def _start_show(self):
111 sx, sy = self._terminal_size
112 self.print(ansi_move_to % (sy-self._height+1, 1))
114 def _end_show(self):
115 sx, sy = self._terminal_size
116 self.print(ansi_move_to % (sy-self._height, 1))
117 self.print(ansi_clear_right)
119 def _resize(self, height):
120 sx, sy = self._terminal_size
121 k = height - self._height
122 if k > 0:
123 self.print(ansi_scroll_up * k)
124 self.print(ansi_window % (1, sy-height))
125 if k < 0:
126 self.print(ansi_window % (1, sy-height))
127 self.print(ansi_scroll_down * abs(k))
129 self._height = height
131 def draw(self, lines):
132 if self._state == 0:
133 self.start()
135 if self._state != 1:
136 return
138 self._terminal_size = get_terminal_size()
139 sx, sy = self._terminal_size
140 nlines = len(lines)
141 self._resize(nlines)
142 self._start_show()
144 for iline, line in enumerate(reversed(lines)):
145 if len(line) > sx - 1:
146 line = line[:sx-1]
148 self.print(ansi_clear_right + line)
149 if iline != nlines - 1:
150 self.print(ansi_next_line)
152 self._end_show()
153 self.flush()
156class LogStatusViewer(StatusViewer):
158 def draw(self, lines):
159 if lines:
160 logger.info(
161 'Progress:\n%s' % '\n'.join(' '+line for line in lines))
164class DummyStatusViewer(StatusViewer):
165 pass
168class Task(object):
169 def __init__(
170 self, progress, id, name, n, state='working', logger=None,
171 group=None):
173 self._id = id
174 self._name = name
175 self._condition = ''
176 self._ispin = 0
177 self._i = None
178 self._n = n
179 self._done = False
180 assert state in ('waiting', 'working')
181 self._state = state
182 self._progress = progress
183 self._logger = logger
184 self._tcreate = time.time()
185 self._group = group
187 def __enter__(self):
188 return self
190 def __exit__(self, type, value, tb):
191 if type is None:
192 self.done()
193 else:
194 self.fail()
196 def __call__(self, it):
197 try:
198 self._n = len(it)
199 except TypeError:
200 self._n = None
202 clean = False
203 try:
204 n = 0
205 for obj in it:
206 self.update(n)
207 yield obj
208 n += 1
210 self.update(n)
211 clean = True
213 finally:
214 if clean:
215 self.done()
216 else:
217 self.fail()
219 def log(self, s):
220 if self._logger is not None:
221 self._logger.info(s)
223 def get_group_time_start(self):
224 if self._group:
225 return self._group.get_group_time_start()
226 else:
227 return self._tcreate
229 def task(self, *args, **kwargs):
230 kwargs['group'] = self
231 return self._progress.task(*args, **kwargs)
233 def update(self, i=None, condition=''):
234 self._state = 'working'
236 self._condition = condition
238 if i is not None:
239 if self._n is not None:
240 i = min(i, self._n)
242 self._i = i
244 self._progress._update()
246 def done(self, condition=''):
247 self.duration = time.time() - self._tcreate
249 if self._state in ('done', 'failed'):
250 return
252 self._condition = condition
253 self._state = 'done'
254 self._progress._end(self)
255 self.log(str(self))
257 def fail(self, condition=''):
258 self.duration = time.time() - self._tcreate
260 if self._state in ('done', 'failed'):
261 return
263 self._condition = condition
264 self._state = 'failed'
265 self._progress._end(self)
266 self.log(str(self))
268 def _str_state(self):
269 s = self._state
270 if s == 'waiting':
271 return ' '
272 elif s == 'working':
273 self._ispin += 1
274 return spinner[self._ispin % len(spinner)] + ' '
275 elif s == 'done':
276 return symbol_done + ' '
277 elif s == 'failed':
278 return symbol_failed + ' '
279 else:
280 return '? '
282 def _idisplay(self):
283 i = self._i
284 if self._n is not None and i > self._n:
285 i = self._n
286 return i
288 def _str_progress(self):
289 if self._i is None:
290 return self._state
291 elif self._n is None:
292 if self._state != 'working':
293 return '... %s (%i)' % (self._state, self._idisplay())
294 else:
295 return '%i' % self._idisplay()
296 else:
297 if self._state == 'working':
298 nw = len(str(self._n))
299 return (('%' + str(nw) + 'i / %i') % (
300 self._idisplay(), self._n)).center(11)
302 elif self._state == 'failed':
303 return '... %s (%i / %i)' % (
304 self._state, self._idisplay(), self._n)
305 else:
306 return '... %s (%i)' % (self._state, self._n)
308 def _str_percent(self):
309 if self._state == 'working' and self._n is not None and self._n >= 4 \
310 and self._i is not None:
311 return '%3.0f%%' % ((100. * self._i) / self._n)
312 else:
313 return ''
315 def _str_condition(self):
316 if self._condition:
317 return '%s' % self._condition
318 else:
319 return ''
321 def _str_bar(self):
322 if self._state == 'working' and self._n is not None and self._n >= 4 \
323 and self._i is not None:
324 nb = 20
325 fb = nb * float(self._i) / self._n
326 ib = int(fb)
327 ip = int((fb - ib) * (len(blocks)-1))
328 if ib == 0 and ip == 0:
329 ip = 1 # indication of start
330 s = blocks[0] * ib
331 if ib < nb:
332 s += blocks[-1-ip] + (nb - ib - 1) * blocks[-1] + blocks[-2]
334 # s = ' ' + bar[0] + bar[1] * ib + bar[2] * (nb - ib) + bar[3]
335 return s
336 else:
337 return ''
339 def __str__(self):
340 return '%s%-23s %-11s %s%-4s %s' % (
341 self._str_state(),
342 self._name,
343 self._str_progress(),
344 self._str_bar(),
345 self._str_percent(),
346 self._str_condition())
349class Progress(object):
351 def __init__(self):
352 self._current_id = 0
353 self._current_group_id = 0
354 self._tasks = {}
355 self._tasks_done = []
356 self._last_update = 0.0
357 self._terms = []
359 def view(self, viewer=None):
360 if g_force_viewer_off or self._terms:
361 viewer = 'off'
362 elif viewer is None:
363 viewer = g_viewer
365 try:
366 term = g_viewer_classes[viewer](self)
367 except KeyError:
368 raise ValueError('Invalid viewer choice: %s' % viewer)
370 self._terms.append(term)
371 return term
373 def hide(self, term):
374 self._update(force=True)
375 self._terms.remove(term)
377 def task(self, name, n=None, logger=None, group=None):
378 self._current_id += 1
379 task = Task(
380 self, self._current_id, name, n, logger=logger, group=group)
381 self._tasks[task._id] = task
382 self._update(force=True)
383 return task
385 def _end(self, task):
386 del self._tasks[task._id]
387 self._tasks_done.append(task)
388 self._update(force=True)
390 def _update(self, force=False):
391 now = time.time()
392 if self._last_update + 0.1 < now or force:
393 self._tasks_done = []
395 lines = self._lines()
396 for term in self._terms:
397 term.draw(lines)
399 self._last_update = now
401 def _lines(self):
402 task_ids = sorted(self._tasks)
403 lines = []
404 for task_id in task_ids:
405 task = self._tasks[task_id]
406 lines.extend(str(task).splitlines())
408 return lines
411g_viewer_classes = {
412 'terminal': TerminalStatusViewer,
413 'log': LogStatusViewer,
414 'off': DummyStatusViewer}
416progress = Progress()