1from __future__ import absolute_import, print_function
3import sys
4import time
6from .get_terminal_size import get_terminal_size
9# spinner = u'\u25dc\u25dd\u25de\u25df'
10# spinner = '⣾⣽⣻⢿⡿⣟⣯⣷'
11spinner = "◴◷◶◵"
12skull = u'\u2620'
13check = u'\u2714'
14cross = u'\u2716'
15bar = u'[- ]'
16blocks = u'\u2588\u2589\u258a\u258b\u258c\u258d\u258e\u258f '
18symbol_done = check
19symbol_failed = cross # skull
21ansi_up = u'\033[%iA'
22ansi_down = u'\033[%iB'
23ansi_left = u'\033[%iC'
24ansi_right = u'\033[%iD'
25ansi_next_line = u'\033E'
27ansi_erase_display = u'\033[2J'
28ansi_window = u'\033[%i;%ir'
29ansi_move_to = u'\033[%i;%iH'
31ansi_clear_down = u'\033[0J'
32ansi_clear_up = u'\033[1J'
33ansi_clear = u'\033[2J'
35ansi_clear_right = u'\033[0K'
37ansi_scroll_up = u'\033D'
38ansi_scroll_down = u'\033M'
40ansi_reset = u'\033c'
43g_force_viewer_off = False
46class TerminalStatusWindow(object):
47 def __init__(self, parent=None):
48 self._terminal_size = get_terminal_size()
49 self._height = 0
50 self._state = 0
51 self._parent = parent
53 def __enter__(self):
54 return self
56 def __exit__(self, *_):
57 self.stop()
59 def print(self, s):
60 print(s, end='', file=sys.stderr)
62 def flush(self):
63 print('', end='', flush=True, file=sys.stderr)
65 def start(self):
66 sx, sy = self._terminal_size
67 self._state = 1
69 def stop(self):
70 if self._state == 1:
71 sx, sy = self._terminal_size
72 self._resize(0)
73 self.print(ansi_move_to % (sy-self._height, 1))
74 self.flush()
76 self._state = 2
77 if self._parent:
78 self._parent.hide(self)
80 def _start_show(self):
81 sx, sy = self._terminal_size
82 self.print(ansi_move_to % (sy-self._height+1, 1))
84 def _end_show(self):
85 sx, sy = self._terminal_size
86 self.print(ansi_move_to % (sy-self._height, 1))
87 self.print(ansi_clear_right)
89 def _resize(self, height):
90 sx, sy = self._terminal_size
91 k = height - self._height
92 if k > 0:
93 self.print(ansi_scroll_up * k)
94 self.print(ansi_window % (1, sy-height))
95 if k < 0:
96 self.print(ansi_window % (1, sy-height))
97 self.print(ansi_scroll_down * abs(k))
99 self._height = height
101 def draw(self, lines):
102 if self._state == 0:
103 self.start()
105 if self._state != 1:
106 return
108 self._terminal_size = get_terminal_size()
109 sx, sy = self._terminal_size
110 nlines = len(lines)
111 self._resize(nlines)
112 self._start_show()
114 for iline, line in enumerate(reversed(lines)):
115 if len(line) > sx - 1:
116 line = line[:sx-1]
118 self.print(ansi_clear_right + line)
119 if iline != nlines - 1:
120 self.print(ansi_next_line)
122 self._end_show()
123 self.flush()
126class DummyStatusWindow(object):
128 def __init__(self, parent=None):
129 self._parent = parent
131 def __enter__(self):
132 return self
134 def __exit__(self, *_):
135 self.stop()
137 def stop(self):
138 if self._parent:
139 self._parent.hide(self)
141 def draw(self, lines):
142 pass
145class Task(object):
146 def __init__(
147 self, progress, id, name, n, state='working', logger=None,
148 group=None):
150 self._id = id
151 self._name = name
152 self._condition = ''
153 self._ispin = 0
154 self._i = None
155 self._n = n
156 self._done = False
157 assert state in ('waiting', 'working')
158 self._state = state
159 self._progress = progress
160 self._logger = logger
161 self._tcreate = time.time()
162 self._group = group
164 def __enter__(self):
165 return self
167 def __exit__(self, type, value, tb):
168 if type is None:
169 self.done()
170 else:
171 self.fail()
173 def __call__(self, it):
174 try:
175 self._n = len(it)
176 except TypeError:
177 self._n = None
179 clean = False
180 try:
181 n = 0
182 for obj in it:
183 self.update(n)
184 yield obj
185 n += 1
187 self.update(n)
188 clean = True
190 finally:
191 if clean:
192 self.done()
193 else:
194 self.fail()
196 def log(self, s):
197 if self._logger is not None:
198 self._logger.info(s)
200 def get_group_time_start(self):
201 if self._group:
202 return self._group.get_group_time_start()
203 else:
204 return self._tcreate
206 def task(self, *args, **kwargs):
207 kwargs['group'] = self
208 return self._progress.task(*args, **kwargs)
210 def update(self, i=None, condition=''):
211 self._state = 'working'
213 self._condition = condition
215 if i is not None:
216 if self._n is not None:
217 i = min(i, self._n)
219 self._i = i
221 self._progress._update()
223 def done(self, condition=''):
224 self.duration = time.time() - self._tcreate
226 if self._state in ('done', 'failed'):
227 return
229 self._condition = condition
230 self._state = 'done'
231 self._progress._end(self)
232 self.log(str(self))
234 def fail(self, condition=''):
235 self.duration = time.time() - self._tcreate
237 if self._state in ('done', 'failed'):
238 return
240 self._condition = condition
241 self._state = 'failed'
242 self._progress._end(self)
243 self.log(str(self))
245 def _str_state(self):
246 s = self._state
247 if s == 'waiting':
248 return ' '
249 elif s == 'working':
250 self._ispin += 1
251 return spinner[self._ispin % len(spinner)] + ' '
252 elif s == 'done':
253 return symbol_done + ' '
254 elif s == 'failed':
255 return symbol_failed + ' '
256 else:
257 return '? '
259 def _idisplay(self):
260 i = self._i
261 if self._n is not None and i > self._n:
262 i = self._n
263 return i
265 def _str_progress(self):
266 if self._i is None:
267 return self._state
268 elif self._n is None:
269 if self._state != 'working':
270 return '... %s (%i)' % (self._state, self._idisplay())
271 else:
272 return '%i' % self._idisplay()
273 else:
274 if self._state == 'working':
275 nw = len(str(self._n))
276 return (('%' + str(nw) + 'i / %i') % (
277 self._idisplay(), self._n)).center(11)
279 elif self._state == 'failed':
280 return '... %s (%i / %i)' % (
281 self._state, self._idisplay(), self._n)
282 else:
283 return '... %s (%i)' % (self._state, self._n)
285 def _str_percent(self):
286 if self._state == 'working' and self._n is not None and self._n >= 4 \
287 and self._i is not None:
288 return '%3.0f%%' % ((100. * self._i) / self._n)
289 else:
290 return ''
292 def _str_condition(self):
293 if self._condition:
294 return '%s' % self._condition
295 else:
296 return ''
298 def _str_bar(self):
299 if self._state == 'working' and self._n is not None and self._n >= 4 \
300 and self._i is not None:
301 nb = 20
302 fb = nb * float(self._i) / self._n
303 ib = int(fb)
304 ip = int((fb - ib) * (len(blocks)-1))
305 if ib == 0 and ip == 0:
306 ip = 1 # indication of start
307 s = blocks[0] * ib
308 if ib < nb:
309 s += blocks[-1-ip] + (nb - ib - 1) * blocks[-1] + blocks[-2]
311 # s = ' ' + bar[0] + bar[1] * ib + bar[2] * (nb - ib) + bar[3]
312 return s
313 else:
314 return ''
316 def __str__(self):
317 return '%s%-23s %-11s %s%-4s %s' % (
318 self._str_state(),
319 self._name,
320 self._str_progress(),
321 self._str_bar(),
322 self._str_percent(),
323 self._str_condition())
326class Progress(object):
328 def __init__(self):
329 self._current_id = 0
330 self._current_group_id = 0
331 self._tasks = {}
332 self._tasks_done = []
333 self._last_update = 0.0
334 self._terms = []
336 def view(self, viewer='terminal'):
337 if g_force_viewer_off or self._terms:
338 viewer = 'off'
340 if viewer == 'terminal':
341 term = TerminalStatusWindow(self)
342 elif viewer == 'off':
343 term = DummyStatusWindow(self)
344 else:
345 raise ValueError('Invalid viewer choice: %s' % viewer)
347 self._terms.append(term)
348 return term
350 def hide(self, term):
351 self._update(force=True)
352 self._terms.remove(term)
354 def task(self, name, n=None, logger=None, group=None):
355 self._current_id += 1
356 task = Task(
357 self, self._current_id, name, n, logger=logger, group=group)
358 self._tasks[task._id] = task
359 self._update(force=True)
360 return task
362 def _end(self, task):
363 del self._tasks[task._id]
364 self._tasks_done.append(task)
365 self._update(force=True)
367 def _update(self, force=False):
368 now = time.time()
369 if self._last_update + 0.1 < now or force:
370 self._tasks_done = []
372 lines = self._lines()
373 for term in self._terms:
374 term.draw(lines)
376 self._last_update = now
378 def _lines(self):
379 task_ids = sorted(self._tasks)
380 lines = []
381 for task_id in task_ids:
382 task = self._tasks[task_id]
383 lines.extend(str(task).splitlines())
385 return lines
388progress = Progress()