Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/progress.py: 74%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Inform users about the progress and success/fail state of long-running tasks. 

8''' 

9 

10import sys 

11import time 

12import logging 

13import threading 

14 

15from shutil import get_terminal_size 

16 

17logger = logging.getLogger('pyrocko.progress') 

18 

19g_spinners = [ 

20 '⣾⣽⣻⢿⡿⣟⣯⣷', 

21 '◴◷◶◵', 

22 '\u25dc\u25dd\u25de\u25df', 

23 '0123456789'] 

24 

25 

26skull = u'\u2620' 

27check = u'\u2714' 

28cross = u'\u2716' 

29bar = u'[- ]' 

30blocks = u'\u2588\u2589\u258a\u258b\u258c\u258d\u258e\u258f ' 

31 

32symbol_done = check 

33symbol_failed = cross # skull 

34 

35ansi_up = u'\033[%iA' 

36ansi_down = u'\033[%iB' 

37ansi_left = u'\033[%iC' 

38ansi_right = u'\033[%iD' 

39ansi_next_line = u'\033E' 

40ansi_previous_line = u'\033[1F' 

41 

42ansi_erase_display = u'\033[2J' 

43ansi_window = u'\033[%i;%ir' 

44ansi_move_to = u'\033[%i;%iH' 

45 

46ansi_clear_down = u'\033[0J' 

47ansi_clear_up = u'\033[1J' 

48ansi_clear = u'\033[2J' 

49 

50ansi_clear_right = u'\033[0K' 

51 

52ansi_scroll_up = u'\033D' 

53ansi_scroll_down = u'\033M' 

54 

55ansi_reset = u'\033c' 

56 

57ansi_save = u'\033 7' 

58ansi_restore = u'\033 8' 

59 

60 

61g_force_viewer_off = False 

62 

63g_viewer = 'terminal' 

64 

65 

66def set_default_viewer(viewer): 

67 ''' 

68 Set default viewer for progress indicators. 

69 

70 :param viewer: 

71 Name of viewer, choices: ``'terminal'``, ``'log'``, ``'off'``, default: 

72 ``'terminal'``. 

73 :type viewer: 

74 str 

75 ''' 

76 

77 global g_viewer 

78 assert viewer in g_viewer_classes 

79 g_viewer = viewer 

80 

81 

82class StatusViewer(object): 

83 

84 def __init__(self, parent, interval=0.1, delay=1.0): 

85 self._parent = parent 

86 self._interval = interval 

87 self._delay = delay 

88 self._last_update = 0.0 

89 self._created = time.time() 

90 

91 def __enter__(self): 

92 return self 

93 

94 def __exit__(self, *_): 

95 self.cleanup() 

96 

97 if self._parent: 

98 self._parent._remove_viewer(self) 

99 

100 def cleanup(self): 

101 pass 

102 

103 def update(self, force): 

104 now = time.time() 

105 if now < self._created + self._delay: 

106 return 

107 

108 if self._last_update + self._interval < now or force: 

109 self._draw() 

110 self._last_update = now 

111 

112 def _draw(self): 

113 pass 

114 

115 

116class DummyStatusViewer(StatusViewer): 

117 pass 

118 

119 

120class LogStatusViewer(StatusViewer): 

121 

122 def __init__(self, parent): 

123 StatusViewer.__init__(self, parent, delay=5., interval=5.) 

124 

125 def _draw(self): 

126 lines = self._parent._render('log') 

127 if lines: 

128 logger.info( 

129 'Progress:\n%s' % '\n'.join(' '+line for line in lines)) 

130 

131 

132class TerminalStatusViewer(StatusViewer): 

133 def __init__(self, parent): 

134 StatusViewer.__init__(self, parent) 

135 self._terminal_size = get_terminal_size() 

136 self._height = 0 

137 self._state = 0 

138 self._nlines_max = 0 

139 self._isatty = sys.stdout.isatty() 

140 

141 def cleanup(self): 

142 if self._state == 1: 

143 sx, sy = self._terminal_size 

144 self._reset() 

145 self._flush() 

146 

147 self._state = 2 

148 

149 def _draw(self): 

150 lines = self._parent._render('terminal') 

151 if self._state == 0: 

152 self._state = 1 

153 

154 if self._state != 1: 

155 return 

156 

157 self._terminal_size = get_terminal_size() 

158 sx, sy = self._terminal_size 

159 nlines = len(lines) 

160 if self._nlines_max < nlines: 

161 self._nlines_max = nlines 

162 self._resize(self._nlines_max) 

163 

164 self._start_show() 

165 

166 for i in range(self._nlines_max - nlines): 

167 lines.append('') 

168 

169 for iline, line in enumerate(reversed(lines)): 

170 if len(line) > sx - 1: 

171 line = line[:sx-1] 

172 

173 self._print(ansi_clear_right + line) 

174 if iline != self._nlines_max - 1: 

175 self._print(ansi_next_line) 

176 

177 self._end_show() 

178 self._flush() 

179 

180 def _print(self, s): 

181 if self._isatty: 

182 print(s, end='', file=sys.stderr) 

183 

184 def _flush(self): 

185 if self._isatty: 

186 print('', end='', flush=True, file=sys.stderr) 

187 

188 def _reset(self): 

189 sx, sy = self._terminal_size 

190 self._print(ansi_window % (1, sy)) 

191 self._print(ansi_move_to % (sy-self._height, 1)) 

192 self._print(ansi_clear_down) 

193 self._height = 0 

194 

195 def _resize(self, height): 

196 sx, sy = self._terminal_size 

197 k = height - self._height 

198 if k > 0: 

199 self._print(ansi_scroll_up * k) 

200 self._print(ansi_window % (1, sy-height)) 

201 if k < 0: 

202 self._print(ansi_window % (1, sy-height)) 

203 self._print(ansi_scroll_down * abs(k)) 

204 

205 self._height = height 

206 

207 def _start_show(self): 

208 sx, sy = self._terminal_size 

209 self._print(ansi_move_to % (sy-self._height+1, 1)) 

210 

211 def _end_show(self): 

212 sx, sy = self._terminal_size 

213 self._print(ansi_move_to % (sy-self._height, 1)) 

214 self._print(ansi_clear_right) 

215 

216 

217class Task(object): 

218 def __init__( 

219 self, progress, id, name, n, state='working', logger=None, 

220 group=None, spinner=g_spinners[0]): 

221 

222 self._id = id 

223 self._name = name 

224 self._condition = '' 

225 self._ispin0 = 0 

226 self._ispin0_last = 0 

227 self._ispin = 0 

228 self._spinner = spinner 

229 self._i = None 

230 self._n = n 

231 self._done = False 

232 assert state in ('waiting', 'working') 

233 self._state = state 

234 self._progress = progress 

235 self._logger = logger 

236 self._tcreate = time.time() 

237 self._group = group 

238 self._lock = threading.RLock() 

239 

240 def __enter__(self): 

241 return self 

242 

243 def __exit__(self, type, value, tb): 

244 if type is None: 

245 self.done() 

246 else: 

247 self.fail() 

248 

249 def __call__(self, it): 

250 try: 

251 self._n = len(it) 

252 except TypeError: 

253 self._n = None 

254 

255 clean = False 

256 try: 

257 n = 0 

258 for obj in it: 

259 self.update(n) 

260 yield obj 

261 n += 1 

262 

263 self.update(n) 

264 clean = True 

265 

266 finally: 

267 if clean: 

268 self.done() 

269 else: 

270 self.fail() 

271 

272 def task(self, *args, **kwargs): 

273 kwargs['group'] = self 

274 return self._progress.task(*args, **kwargs) 

275 

276 def update(self, i=None, condition=''): 

277 with self._lock: 

278 self._ispin0 += 1 

279 self._state = 'working' 

280 self._condition = condition 

281 if i is not None: 

282 if self._n is not None: 

283 i = min(i, self._n) 

284 

285 self._i = i 

286 

287 self._progress._update(False) 

288 

289 def done(self, condition=''): 

290 with self._lock: 

291 self.duration = time.time() - self._tcreate 

292 

293 if self._state in ('done', 'failed'): 

294 return 

295 

296 self._condition = condition 

297 self._state = 'done' 

298 self._log(str(self)) 

299 self._progress._task_end(self) 

300 

301 def fail(self, condition=''): 

302 with self._lock: 

303 self.duration = time.time() - self._tcreate 

304 

305 if self._state in ('done', 'failed'): 

306 return 

307 

308 self._condition = condition 

309 self._state = 'failed' 

310 self._log(str(self)) 

311 self._progress._task_end(self) 

312 

313 def _log(self, s): 

314 if self._logger is not None: 

315 self._logger.debug(s) 

316 

317 def _get_group_time_start(self): 

318 if self._group: 

319 return self._group._get_group_time_start() 

320 else: 

321 return self._tcreate 

322 

323 def _str_state(self): 

324 s = self._state 

325 if s == 'waiting': 

326 return ' ' 

327 elif s == 'working': 

328 if self._ispin0_last != self._ispin0: 

329 self._ispin += 1 

330 self._ispin0_last = self._ispin0 

331 

332 return self._spinner[self._ispin % len(self._spinner)] + ' ' 

333 elif s == 'done': 

334 return symbol_done + ' ' 

335 elif s == 'failed': 

336 return symbol_failed + ' ' 

337 else: 

338 return '? ' 

339 

340 def _idisplay(self): 

341 i = self._i 

342 if self._n is not None and i > self._n: 

343 i = self._n 

344 return i 

345 

346 def _str_progress(self): 

347 if self._i is None: 

348 return self._state 

349 elif self._n is None: 

350 if self._state != 'working': 

351 return '... %s (%i)' % (self._state, self._idisplay()) 

352 else: 

353 return '%i' % self._idisplay() 

354 else: 

355 if self._state == 'working': 

356 nw = len(str(self._n)) 

357 return (('%' + str(nw) + 'i / %i') % ( 

358 self._idisplay(), self._n)).center(11) 

359 

360 elif self._state == 'failed': 

361 return '... %s (%i / %i)' % ( 

362 self._state, self._idisplay(), self._n) 

363 else: 

364 return '... %s (%i)' % (self._state, self._n) 

365 

366 def _str_percent(self): 

367 if self._state == 'working' and self._n is not None and self._n >= 4 \ 

368 and self._i is not None: 

369 return '%3.0f%%' % ((100. * self._i) / self._n) 

370 else: 

371 return '' 

372 

373 def _str_condition(self): 

374 if self._condition: 

375 return '%s' % self._condition 

376 else: 

377 return '' 

378 

379 def _str_bar(self): 

380 if self._state == 'working' and self._n is not None and self._n >= 4 \ 

381 and self._i is not None: 

382 nb = 20 

383 fb = nb * float(self._i) / self._n 

384 ib = int(fb) 

385 ip = int((fb - ib) * (len(blocks)-1)) 

386 if ib == 0 and ip == 0: 

387 ip = 1 # indication of start 

388 s = blocks[0] * ib 

389 if ib < nb: 

390 s += blocks[-1-ip] + (nb - ib - 1) * blocks[-1] + blocks[-2] 

391 

392 # s = ' ' + bar[0] + bar[1] * ib + bar[2] * (nb - ib) + bar[3] 

393 return s 

394 else: 

395 return '' 

396 

397 def _render(self, style): 

398 if style == 'terminal': 

399 return '%s%-23s %-11s %s%-4s %s' % ( 

400 self._str_state(), 

401 self._name, 

402 self._str_progress(), 

403 self._str_bar(), 

404 self._str_percent(), 

405 self._str_condition()) 

406 

407 elif style == 'log': 

408 return '%s: %-23s %s%-4s %s' % ( 

409 self._state, 

410 self._name, 

411 self._str_progress(), 

412 self._str_percent(), 

413 self._str_condition()) 

414 else: 

415 return '' 

416 

417 def __str__(self): 

418 return '%s%-23s %s%-4s %s' % ( 

419 self._str_state(), 

420 self._name, 

421 self._str_progress(), 

422 self._str_percent(), 

423 self._str_condition()) 

424 

425 

426class Progress(object): 

427 

428 def __init__(self): 

429 self._current_id = 0 

430 self._tasks = {} 

431 self._viewers = [] 

432 self._lock = threading.RLock() 

433 self._isatty = sys.stdout.isatty() 

434 

435 def view(self, viewer=None): 

436 if g_force_viewer_off or self._viewers: 

437 viewer = 'off' 

438 elif viewer is None: 

439 viewer = g_viewer 

440 

441 if not self._isatty and viewer == 'terminal': 

442 logger.debug('No tty attached, switching to log progress viewer.') 

443 viewer = 'log' 

444 

445 try: 

446 viewer = g_viewer_classes[viewer](self) 

447 except KeyError: 

448 raise ValueError('Invalid viewer choice: %s' % viewer) 

449 

450 self._viewers.append(viewer) 

451 return viewer 

452 

453 def _remove_viewer(self, viewer): 

454 self._update(True) 

455 self._viewers.remove(viewer) 

456 

457 def task(self, name, n=None, logger=None, group=None): 

458 with self._lock: 

459 self._current_id += 1 

460 task = Task( 

461 self, self._current_id, name, n, logger=logger, group=group) 

462 self._tasks[task._id] = task 

463 self._update(True) 

464 return task 

465 

466 def _task_end(self, task): 

467 with self._lock: 

468 self._update(True) 

469 del self._tasks[task._id] 

470 self._update(True) 

471 

472 def _update(self, force): 

473 with self._lock: 

474 for viewer in self._viewers: 

475 viewer.update(force) 

476 

477 def _render(self, style): 

478 task_ids = sorted(self._tasks) 

479 lines = [] 

480 for task_id in task_ids: 

481 task = self._tasks[task_id] 

482 lines.extend(task._render(style).splitlines()) 

483 

484 return lines 

485 

486 def _debug_log(self): 

487 logger.debug( 

488 'Viewers: %s\n Tasks active: %i\n Tasks total: %i', 

489 ', '.join(viewer.__class__.__name__ for viewer in self._viewers) 

490 if self._viewers else 'none', 

491 len(self._tasks), 

492 self._current_id) 

493 

494 

495g_viewer_classes = { 

496 'terminal': TerminalStatusViewer, 

497 'log': LogStatusViewer, 

498 'off': DummyStatusViewer} 

499 

500g_progress = Progress() 

501progress = g_progress # compatibility 

502view = g_progress.view 

503task = g_progress.task