1from __future__ import absolute_import, print_function
3import sys
4import time
5import logging
7from .get_terminal_size import get_terminal_size
9logger = logging.getLogger('pyrocko.progress')
11# TODO: Refactor so that if multiple viewers are attached, they can do
12# their updates independently of each other (at independent time intervals).
14# TODO: Refactor so that different viewers can render task states differently.
16# spinner = u'\u25dc\u25dd\u25de\u25df'
17# spinner = '⣾⣽⣻⢿⡿⣟⣯⣷'
18spinner = "◴◷◶◵"
19skull = u'\u2620'
20check = u'\u2714'
21cross = u'\u2716'
22bar = u'[- ]'
23blocks = u'\u2588\u2589\u258a\u258b\u258c\u258d\u258e\u258f '
25symbol_done = check
26symbol_failed = cross # skull
28ansi_up = u'\033[%iA'
29ansi_down = u'\033[%iB'
30ansi_left = u'\033[%iC'
31ansi_right = u'\033[%iD'
32ansi_next_line = u'\033E'
34ansi_erase_display = u'\033[2J'
35ansi_window = u'\033[%i;%ir'
36ansi_move_to = u'\033[%i;%iH'
38ansi_clear_down = u'\033[0J'
39ansi_clear_up = u'\033[1J'
40ansi_clear = u'\033[2J'
42ansi_clear_right = u'\033[0K'
44ansi_scroll_up = u'\033D'
45ansi_scroll_down = u'\033M'
47ansi_reset = u'\033c'
50g_force_viewer_off = False
52g_viewer = 'terminal'
55def set_default_viewer(viewer):
56 global g_viewer
57 g_viewer = viewer
60class StatusViewer(object):
62 def __init__(self, parent=None):
63 self._parent = parent
65 def __enter__(self):
66 return self
68 def __exit__(self, *_):
69 self.stop()
71 def stop(self):
72 if self._parent:
73 self._parent.hide(self)
75 def draw(self, lines):
76 pass
79class TerminalStatusViewer(StatusViewer):
80 def __init__(self, parent=None):
81 self._terminal_size = get_terminal_size()
82 self._height = 0
83 self._state = 0
84 self._parent = parent
86 def print(self, s):
87 print(s, end='', file=sys.stderr)
89 def flush(self):
90 print('', end='', flush=True, file=sys.stderr)
92 def start(self):
93 sx, sy = self._terminal_size
94 self._state = 1
96 def stop(self):
97 if self._state == 1:
98 sx, sy = self._terminal_size
99 self._resize(0)
100 self.print(ansi_move_to % (sy-self._height, 1))
101 self.flush()
103 self._state = 2
104 if self._parent:
105 self._parent.hide(self)
107 def _start_show(self):
108 sx, sy = self._terminal_size
109 self.print(ansi_move_to % (sy-self._height+1, 1))
111 def _end_show(self):
112 sx, sy = self._terminal_size
113 self.print(ansi_move_to % (sy-self._height, 1))
114 self.print(ansi_clear_right)
116 def _resize(self, height):
117 sx, sy = self._terminal_size
118 k = height - self._height
119 if k > 0:
120 self.print(ansi_scroll_up * k)
121 self.print(ansi_window % (1, sy-height))
122 if k < 0:
123 self.print(ansi_window % (1, sy-height))
124 self.print(ansi_scroll_down * abs(k))
126 self._height = height
128 def draw(self, lines):
129 if self._state == 0:
130 self.start()
132 if self._state != 1:
133 return
135 self._terminal_size = get_terminal_size()
136 sx, sy = self._terminal_size
137 nlines = len(lines)
138 self._resize(nlines)
139 self._start_show()
141 for iline, line in enumerate(reversed(lines)):
142 if len(line) > sx - 1:
143 line = line[:sx-1]
145 self.print(ansi_clear_right + line)
146 if iline != nlines - 1:
147 self.print(ansi_next_line)
149 self._end_show()
150 self.flush()
153class LogStatusViewer(StatusViewer):
155 def draw(self, lines):
156 if lines:
157 logger.info(
158 'Progress:\n%s' % '\n'.join(' '+line for line in lines))
161class DummyStatusViewer(StatusViewer):
162 pass
165class Task(object):
166 def __init__(
167 self, progress, id, name, n, state='working', logger=None,
168 group=None):
170 self._id = id
171 self._name = name
172 self._condition = ''
173 self._ispin = 0
174 self._i = None
175 self._n = n
176 self._done = False
177 assert state in ('waiting', 'working')
178 self._state = state
179 self._progress = progress
180 self._logger = logger
181 self._tcreate = time.time()
182 self._group = group
184 def __enter__(self):
185 return self
187 def __exit__(self, type, value, tb):
188 if type is None:
189 self.done()
190 else:
191 self.fail()
193 def __call__(self, it):
194 try:
195 self._n = len(it)
196 except TypeError:
197 self._n = None
199 clean = False
200 try:
201 n = 0
202 for obj in it:
203 self.update(n)
204 yield obj
205 n += 1
207 self.update(n)
208 clean = True
210 finally:
211 if clean:
212 self.done()
213 else:
214 self.fail()
216 def log(self, s):
217 if self._logger is not None:
218 self._logger.info(s)
220 def get_group_time_start(self):
221 if self._group:
222 return self._group.get_group_time_start()
223 else:
224 return self._tcreate
226 def task(self, *args, **kwargs):
227 kwargs['group'] = self
228 return self._progress.task(*args, **kwargs)
230 def update(self, i=None, condition=''):
231 self._state = 'working'
233 self._condition = condition
235 if i is not None:
236 if self._n is not None:
237 i = min(i, self._n)
239 self._i = i
241 self._progress._update()
243 def done(self, condition=''):
244 self.duration = time.time() - self._tcreate
246 if self._state in ('done', 'failed'):
247 return
249 self._condition = condition
250 self._state = 'done'
251 self._progress._end(self)
252 self.log(str(self))
254 def fail(self, condition=''):
255 self.duration = time.time() - self._tcreate
257 if self._state in ('done', 'failed'):
258 return
260 self._condition = condition
261 self._state = 'failed'
262 self._progress._end(self)
263 self.log(str(self))
265 def _str_state(self):
266 s = self._state
267 if s == 'waiting':
268 return ' '
269 elif s == 'working':
270 self._ispin += 1
271 return spinner[self._ispin % len(spinner)] + ' '
272 elif s == 'done':
273 return symbol_done + ' '
274 elif s == 'failed':
275 return symbol_failed + ' '
276 else:
277 return '? '
279 def _idisplay(self):
280 i = self._i
281 if self._n is not None and i > self._n:
282 i = self._n
283 return i
285 def _str_progress(self):
286 if self._i is None:
287 return self._state
288 elif self._n is None:
289 if self._state != 'working':
290 return '... %s (%i)' % (self._state, self._idisplay())
291 else:
292 return '%i' % self._idisplay()
293 else:
294 if self._state == 'working':
295 nw = len(str(self._n))
296 return (('%' + str(nw) + 'i / %i') % (
297 self._idisplay(), self._n)).center(11)
299 elif self._state == 'failed':
300 return '... %s (%i / %i)' % (
301 self._state, self._idisplay(), self._n)
302 else:
303 return '... %s (%i)' % (self._state, self._n)
305 def _str_percent(self):
306 if self._state == 'working' and self._n is not None and self._n >= 4 \
307 and self._i is not None:
308 return '%3.0f%%' % ((100. * self._i) / self._n)
309 else:
310 return ''
312 def _str_condition(self):
313 if self._condition:
314 return '%s' % self._condition
315 else:
316 return ''
318 def _str_bar(self):
319 if self._state == 'working' and self._n is not None and self._n >= 4 \
320 and self._i is not None:
321 nb = 20
322 fb = nb * float(self._i) / self._n
323 ib = int(fb)
324 ip = int((fb - ib) * (len(blocks)-1))
325 if ib == 0 and ip == 0:
326 ip = 1 # indication of start
327 s = blocks[0] * ib
328 if ib < nb:
329 s += blocks[-1-ip] + (nb - ib - 1) * blocks[-1] + blocks[-2]
331 # s = ' ' + bar[0] + bar[1] * ib + bar[2] * (nb - ib) + bar[3]
332 return s
333 else:
334 return ''
336 def __str__(self):
337 return '%s%-23s %-11s %s%-4s %s' % (
338 self._str_state(),
339 self._name,
340 self._str_progress(),
341 self._str_bar(),
342 self._str_percent(),
343 self._str_condition())
346class Progress(object):
348 def __init__(self):
349 self._current_id = 0
350 self._current_group_id = 0
351 self._tasks = {}
352 self._tasks_done = []
353 self._last_update = 0.0
354 self._terms = []
356 def view(self, viewer=None):
357 if g_force_viewer_off or self._terms:
358 viewer = 'off'
359 elif viewer is None:
360 viewer = g_viewer
362 try:
363 term = g_viewer_classes[viewer](self)
364 except KeyError:
365 raise ValueError('Invalid viewer choice: %s' % viewer)
367 self._terms.append(term)
368 return term
370 def hide(self, term):
371 self._update(force=True)
372 self._terms.remove(term)
374 def task(self, name, n=None, logger=None, group=None):
375 self._current_id += 1
376 task = Task(
377 self, self._current_id, name, n, logger=logger, group=group)
378 self._tasks[task._id] = task
379 self._update(force=True)
380 return task
382 def _end(self, task):
383 del self._tasks[task._id]
384 self._tasks_done.append(task)
385 self._update(force=True)
387 def _update(self, force=False):
388 now = time.time()
389 if self._last_update + 0.1 < now or force:
390 self._tasks_done = []
392 lines = self._lines()
393 for term in self._terms:
394 term.draw(lines)
396 self._last_update = now
398 def _lines(self):
399 task_ids = sorted(self._tasks)
400 lines = []
401 for task_id in task_ids:
402 task = self._tasks[task_id]
403 lines.extend(str(task).splitlines())
405 return lines
408g_viewer_classes = {
409 'terminal': TerminalStatusViewer,
410 'log': LogStatusViewer,
411 'off': DummyStatusViewer}
413progress = Progress()