Coverage for /usr/local/lib/python3.13/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 self.update(0) 

240 

241 def __enter__(self): 

242 return self 

243 

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

245 if type is None: 

246 self.done() 

247 else: 

248 self.fail() 

249 

250 def __call__(self, it): 

251 try: 

252 self._n = len(it) 

253 except TypeError: 

254 self._n = None 

255 

256 clean = False 

257 try: 

258 n = 0 

259 for obj in it: 

260 self.update(n) 

261 yield obj 

262 n += 1 

263 

264 self.update(n) 

265 clean = True 

266 

267 finally: 

268 if clean: 

269 self.done() 

270 else: 

271 self.fail() 

272 

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

274 kwargs['group'] = self 

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

276 

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

278 with self._lock: 

279 self._ispin0 += 1 

280 self._state = 'working' 

281 self._condition = condition 

282 if i is not None: 

283 if self._n is not None: 

284 i = min(i, self._n) 

285 

286 self._i = i 

287 

288 self._progress._update(False) 

289 

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

291 with self._lock: 

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

293 

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

295 return 

296 

297 self._condition = condition 

298 self._state = 'done' 

299 self._log(str(self)) 

300 self._progress._task_end(self) 

301 

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

303 with self._lock: 

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

305 

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

307 return 

308 

309 self._condition = condition 

310 self._state = 'failed' 

311 self._log(str(self)) 

312 self._progress._task_end(self) 

313 

314 def _log(self, s): 

315 if self._logger is not None: 

316 self._logger.debug(s) 

317 

318 def _get_group_time_start(self): 

319 if self._group: 

320 return self._group._get_group_time_start() 

321 else: 

322 return self._tcreate 

323 

324 def _str_state(self): 

325 s = self._state 

326 if s == 'waiting': 

327 return ' ' 

328 elif s == 'working': 

329 if self._ispin0_last != self._ispin0: 

330 self._ispin += 1 

331 self._ispin0_last = self._ispin0 

332 

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

334 elif s == 'done': 

335 return symbol_done + ' ' 

336 elif s == 'failed': 

337 return symbol_failed + ' ' 

338 else: 

339 return '? ' 

340 

341 def _idisplay(self): 

342 i = self._i 

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

344 i = self._n 

345 return i 

346 

347 def _str_progress(self): 

348 if self._i is None: 

349 return self._state 

350 elif self._n is None: 

351 if self._state != 'working': 

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

353 else: 

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

355 else: 

356 if self._state == 'working': 

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

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

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

360 

361 elif self._state == 'failed': 

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

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

364 else: 

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

366 

367 def _str_percent(self): 

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

369 and self._i is not None: 

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

371 else: 

372 return '' 

373 

374 def _str_condition(self): 

375 if self._condition: 

376 return '%s' % self._condition 

377 else: 

378 return '' 

379 

380 def _str_bar(self): 

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

382 and self._i is not None: 

383 nb = 20 

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

385 ib = int(fb) 

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

387 if ib == 0 and ip == 0: 

388 ip = 1 # indication of start 

389 s = blocks[0] * ib 

390 if ib < nb: 

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

392 

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

394 return s 

395 else: 

396 return '' 

397 

398 def _render(self, style): 

399 if style == 'terminal': 

400 return '%s%-40s %-11s %s%-4s %s' % ( 

401 self._str_state(), 

402 self._name, 

403 self._str_progress(), 

404 self._str_bar(), 

405 self._str_percent(), 

406 self._str_condition()) 

407 

408 elif style == 'log': 

409 return '%s: %-40s %s%-4s %s' % ( 

410 self._state, 

411 self._name, 

412 self._str_progress(), 

413 self._str_percent(), 

414 self._str_condition()) 

415 else: 

416 return '' 

417 

418 def __str__(self): 

419 return '%s%-40s %s%-4s %s' % ( 

420 self._str_state(), 

421 self._name, 

422 self._str_progress(), 

423 self._str_percent(), 

424 self._str_condition()) 

425 

426 

427class Progress(object): 

428 

429 def __init__(self): 

430 self._current_id = 0 

431 self._tasks = {} 

432 self._viewers = [] 

433 self._lock = threading.RLock() 

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

435 

436 def view(self, viewer=None): 

437 if g_force_viewer_off or self._viewers: 

438 viewer = 'off' 

439 elif viewer is None: 

440 viewer = g_viewer 

441 

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

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

444 viewer = 'log' 

445 

446 try: 

447 viewer = g_viewer_classes[viewer](self) 

448 except KeyError: 

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

450 

451 self._viewers.append(viewer) 

452 return viewer 

453 

454 def _remove_viewer(self, viewer): 

455 self._update(True) 

456 self._viewers.remove(viewer) 

457 

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

459 with self._lock: 

460 self._current_id += 1 

461 task = Task( 

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

463 self._tasks[task._id] = task 

464 self._update(True) 

465 return task 

466 

467 def _task_end(self, task): 

468 with self._lock: 

469 self._update(True) 

470 del self._tasks[task._id] 

471 self._update(True) 

472 

473 def _update(self, force): 

474 with self._lock: 

475 for viewer in self._viewers: 

476 viewer.update(force) 

477 

478 def _render(self, style): 

479 task_ids = sorted(self._tasks) 

480 lines = [] 

481 for task_id in task_ids: 

482 task = self._tasks[task_id] 

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

484 

485 return lines 

486 

487 def _debug_log(self): 

488 logger.debug( 

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

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

491 if self._viewers else 'none', 

492 len(self._tasks), 

493 self._current_id) 

494 

495 

496g_viewer_classes = { 

497 'terminal': TerminalStatusViewer, 

498 'log': LogStatusViewer, 

499 'off': DummyStatusViewer} 

500 

501g_progress = Progress() 

502progress = g_progress # compatibility 

503view = g_progress.view 

504task = g_progress.task