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')
12# spinner = u'\u25dc\u25dd\u25de\u25df'
13# spinner = '⣾⣽⣻⢿⡿⣟⣯⣷'
14spinner = "◴◷◶◵"
15skull = u'\u2620'
16check = u'\u2714'
17cross = u'\u2716'
18bar = u'[- ]'
19blocks = u'\u2588\u2589\u258a\u258b\u258c\u258d\u258e\u258f '
21symbol_done = check
22symbol_failed = cross # skull
24ansi_up = u'\033[%iA'
25ansi_down = u'\033[%iB'
26ansi_left = u'\033[%iC'
27ansi_right = u'\033[%iD'
28ansi_next_line = u'\033E'
30ansi_erase_display = u'\033[2J'
31ansi_window = u'\033[%i;%ir'
32ansi_move_to = u'\033[%i;%iH'
34ansi_clear_down = u'\033[0J'
35ansi_clear_up = u'\033[1J'
36ansi_clear = u'\033[2J'
38ansi_clear_right = u'\033[0K'
40ansi_scroll_up = u'\033D'
41ansi_scroll_down = u'\033M'
43ansi_reset = u'\033c'
46g_force_viewer_off = False
48g_viewer = 'terminal'
51def set_default_viewer(viewer):
52 global g_viewer
53 g_viewer = viewer
56class StatusViewer(object):
58 def __init__(self, parent=None):
59 self._parent = parent
61 def __enter__(self):
62 return self
64 def __exit__(self, *_):
65 self.stop()
67 def stop(self):
68 if self._parent:
69 self._parent.hide(self)
71 def draw(self, lines):
72 pass
75class TerminalStatusViewer(StatusViewer):
76 def __init__(self, parent=None):
77 self._terminal_size = get_terminal_size()
78 self._height = 0
79 self._state = 0
80 self._parent = parent
82 def print(self, s):
83 print(s, end='', file=sys.stderr)
85 def flush(self):
86 print('', end='', flush=True, file=sys.stderr)
88 def start(self):
89 sx, sy = self._terminal_size
90 self._state = 1
92 def stop(self):
93 if self._state == 1:
94 sx, sy = self._terminal_size
95 self._resize(0)
96 self.print(ansi_move_to % (sy-self._height, 1))
97 self.flush()
99 self._state = 2
100 if self._parent:
101 self._parent.hide(self)
103 def _start_show(self):
104 sx, sy = self._terminal_size
105 self.print(ansi_move_to % (sy-self._height+1, 1))
107 def _end_show(self):
108 sx, sy = self._terminal_size
109 self.print(ansi_move_to % (sy-self._height, 1))
110 self.print(ansi_clear_right)
112 def _resize(self, height):
113 sx, sy = self._terminal_size
114 k = height - self._height
115 if k > 0:
116 self.print(ansi_scroll_up * k)
117 self.print(ansi_window % (1, sy-height))
118 if k < 0:
119 self.print(ansi_window % (1, sy-height))
120 self.print(ansi_scroll_down * abs(k))
122 self._height = height
124 def draw(self, lines):
125 if self._state == 0:
126 self.start()
128 if self._state != 1:
129 return
131 self._terminal_size = get_terminal_size()
132 sx, sy = self._terminal_size
133 nlines = len(lines)
134 self._resize(nlines)
135 self._start_show()
137 for iline, line in enumerate(reversed(lines)):
138 if len(line) > sx - 1:
139 line = line[:sx-1]
141 self.print(ansi_clear_right + line)
142 if iline != nlines - 1:
143 self.print(ansi_next_line)
145 self._end_show()
146 self.flush()
149class LogStatusViewer(StatusViewer):
151 def draw(self, lines):
152 logger.info('Progress:\n%s' % '\n'.join(' '+line for line in lines))
155class DummyStatusViewer(StatusViewer):
156 pass
159class Task(object):
160 def __init__(
161 self, progress, id, name, n, state='working', logger=None,
162 group=None):
164 self._id = id
165 self._name = name
166 self._condition = ''
167 self._ispin = 0
168 self._i = None
169 self._n = n
170 self._done = False
171 assert state in ('waiting', 'working')
172 self._state = state
173 self._progress = progress
174 self._logger = logger
175 self._tcreate = time.time()
176 self._group = group
178 def __enter__(self):
179 return self
181 def __exit__(self, type, value, tb):
182 if type is None:
183 self.done()
184 else:
185 self.fail()
187 def __call__(self, it):
188 try:
189 self._n = len(it)
190 except TypeError:
191 self._n = None
193 clean = False
194 try:
195 n = 0
196 for obj in it:
197 self.update(n)
198 yield obj
199 n += 1
201 self.update(n)
202 clean = True
204 finally:
205 if clean:
206 self.done()
207 else:
208 self.fail()
210 def log(self, s):
211 if self._logger is not None:
212 self._logger.info(s)
214 def get_group_time_start(self):
215 if self._group:
216 return self._group.get_group_time_start()
217 else:
218 return self._tcreate
220 def task(self, *args, **kwargs):
221 kwargs['group'] = self
222 return self._progress.task(*args, **kwargs)
224 def update(self, i=None, condition=''):
225 self._state = 'working'
227 self._condition = condition
229 if i is not None:
230 if self._n is not None:
231 i = min(i, self._n)
233 self._i = i
235 self._progress._update()
237 def done(self, condition=''):
238 self.duration = time.time() - self._tcreate
240 if self._state in ('done', 'failed'):
241 return
243 self._condition = condition
244 self._state = 'done'
245 self._progress._end(self)
246 self.log(str(self))
248 def fail(self, condition=''):
249 self.duration = time.time() - self._tcreate
251 if self._state in ('done', 'failed'):
252 return
254 self._condition = condition
255 self._state = 'failed'
256 self._progress._end(self)
257 self.log(str(self))
259 def _str_state(self):
260 s = self._state
261 if s == 'waiting':
262 return ' '
263 elif s == 'working':
264 self._ispin += 1
265 return spinner[self._ispin % len(spinner)] + ' '
266 elif s == 'done':
267 return symbol_done + ' '
268 elif s == 'failed':
269 return symbol_failed + ' '
270 else:
271 return '? '
273 def _idisplay(self):
274 i = self._i
275 if self._n is not None and i > self._n:
276 i = self._n
277 return i
279 def _str_progress(self):
280 if self._i is None:
281 return self._state
282 elif self._n is None:
283 if self._state != 'working':
284 return '... %s (%i)' % (self._state, self._idisplay())
285 else:
286 return '%i' % self._idisplay()
287 else:
288 if self._state == 'working':
289 nw = len(str(self._n))
290 return (('%' + str(nw) + 'i / %i') % (
291 self._idisplay(), self._n)).center(11)
293 elif self._state == 'failed':
294 return '... %s (%i / %i)' % (
295 self._state, self._idisplay(), self._n)
296 else:
297 return '... %s (%i)' % (self._state, self._n)
299 def _str_percent(self):
300 if self._state == 'working' and self._n is not None and self._n >= 4 \
301 and self._i is not None:
302 return '%3.0f%%' % ((100. * self._i) / self._n)
303 else:
304 return ''
306 def _str_condition(self):
307 if self._condition:
308 return '%s' % self._condition
309 else:
310 return ''
312 def _str_bar(self):
313 if self._state == 'working' and self._n is not None and self._n >= 4 \
314 and self._i is not None:
315 nb = 20
316 fb = nb * float(self._i) / self._n
317 ib = int(fb)
318 ip = int((fb - ib) * (len(blocks)-1))
319 if ib == 0 and ip == 0:
320 ip = 1 # indication of start
321 s = blocks[0] * ib
322 if ib < nb:
323 s += blocks[-1-ip] + (nb - ib - 1) * blocks[-1] + blocks[-2]
325 # s = ' ' + bar[0] + bar[1] * ib + bar[2] * (nb - ib) + bar[3]
326 return s
327 else:
328 return ''
330 def __str__(self):
331 return '%s%-23s %-11s %s%-4s %s' % (
332 self._str_state(),
333 self._name,
334 self._str_progress(),
335 self._str_bar(),
336 self._str_percent(),
337 self._str_condition())
340class Progress(object):
342 def __init__(self):
343 self._current_id = 0
344 self._current_group_id = 0
345 self._tasks = {}
346 self._tasks_done = []
347 self._last_update = 0.0
348 self._terms = []
350 def view(self, viewer=None):
351 if g_force_viewer_off or self._terms:
352 viewer = 'off'
353 elif viewer is None:
354 viewer = g_viewer
356 try:
357 term = g_viewer_classes[viewer](self)
358 except KeyError:
359 raise ValueError('Invalid viewer choice: %s' % viewer)
361 self._terms.append(term)
362 return term
364 def hide(self, term):
365 self._update(force=True)
366 self._terms.remove(term)
368 def task(self, name, n=None, logger=None, group=None):
369 self._current_id += 1
370 task = Task(
371 self, self._current_id, name, n, logger=logger, group=group)
372 self._tasks[task._id] = task
373 self._update(force=True)
374 return task
376 def _end(self, task):
377 del self._tasks[task._id]
378 self._tasks_done.append(task)
379 self._update(force=True)
381 def _update(self, force=False):
382 now = time.time()
383 if self._last_update + 0.1 < now or force:
384 self._tasks_done = []
386 lines = self._lines()
387 for term in self._terms:
388 term.draw(lines)
390 self._last_update = now
392 def _lines(self):
393 task_ids = sorted(self._tasks)
394 lines = []
395 for task_id in task_ids:
396 task = self._tasks[task_id]
397 lines.extend(str(task).splitlines())
399 return lines
402g_viewer_classes = {
403 'terminal': TerminalStatusViewer,
404 'log': LogStatusViewer,
405 'off': DummyStatusViewer}
408progress = Progress()