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

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 

13 

14from .get_terminal_size import get_terminal_size 

15 

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

17 

18# TODO: Refactor so that if multiple viewers are attached, they can do 

19# their updates independently of each other (at independent time intervals). 

20 

21# TODO: Refactor so that different viewers can render task states differently. 

22 

23# spinner = u'\u25dc\u25dd\u25de\u25df' 

24# spinner = '⣾⣽⣻⢿⡿⣟⣯⣷' 

25spinner = '◴◷◶◵' 

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' 

40 

41ansi_erase_display = u'\033[2J' 

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

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

44 

45ansi_clear_down = u'\033[0J' 

46ansi_clear_up = u'\033[1J' 

47ansi_clear = u'\033[2J' 

48 

49ansi_clear_right = u'\033[0K' 

50 

51ansi_scroll_up = u'\033D' 

52ansi_scroll_down = u'\033M' 

53 

54ansi_reset = u'\033c' 

55 

56 

57g_force_viewer_off = False 

58 

59g_viewer = 'terminal' 

60 

61 

62def set_default_viewer(viewer): 

63 global g_viewer 

64 g_viewer = viewer 

65 

66 

67class StatusViewer(object): 

68 

69 def __init__(self, parent=None): 

70 self._parent = parent 

71 

72 def __enter__(self): 

73 return self 

74 

75 def __exit__(self, *_): 

76 self.stop() 

77 

78 def stop(self): 

79 if self._parent: 

80 self._parent.hide(self) 

81 

82 def draw(self, lines): 

83 pass 

84 

85 

86class TerminalStatusViewer(StatusViewer): 

87 def __init__(self, parent=None): 

88 self._terminal_size = get_terminal_size() 

89 self._height = 0 

90 self._state = 0 

91 self._parent = parent 

92 

93 def print(self, s): 

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

95 

96 def flush(self): 

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

98 

99 def start(self): 

100 sx, sy = self._terminal_size 

101 self._state = 1 

102 

103 def stop(self): 

104 if self._state == 1: 

105 sx, sy = self._terminal_size 

106 self._resize(0) 

107 self.print(ansi_move_to % (sy-self._height, 1)) 

108 self.flush() 

109 

110 self._state = 2 

111 if self._parent: 

112 self._parent.hide(self) 

113 

114 def _start_show(self): 

115 sx, sy = self._terminal_size 

116 self.print(ansi_move_to % (sy-self._height+1, 1)) 

117 

118 def _end_show(self): 

119 sx, sy = self._terminal_size 

120 self.print(ansi_move_to % (sy-self._height, 1)) 

121 self.print(ansi_clear_right) 

122 

123 def _resize(self, height): 

124 sx, sy = self._terminal_size 

125 k = height - self._height 

126 if k > 0: 

127 self.print(ansi_scroll_up * k) 

128 self.print(ansi_window % (1, sy-height)) 

129 if k < 0: 

130 self.print(ansi_window % (1, sy-height)) 

131 self.print(ansi_scroll_down * abs(k)) 

132 

133 self._height = height 

134 

135 def draw(self, lines): 

136 if self._state == 0: 

137 self.start() 

138 

139 if self._state != 1: 

140 return 

141 

142 self._terminal_size = get_terminal_size() 

143 sx, sy = self._terminal_size 

144 nlines = len(lines) 

145 self._resize(nlines) 

146 self._start_show() 

147 

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

149 if len(line) > sx - 1: 

150 line = line[:sx-1] 

151 

152 self.print(ansi_clear_right + line) 

153 if iline != nlines - 1: 

154 self.print(ansi_next_line) 

155 

156 self._end_show() 

157 self.flush() 

158 

159 

160class LogStatusViewer(StatusViewer): 

161 

162 def draw(self, lines): 

163 if lines: 

164 logger.info( 

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

166 

167 

168class DummyStatusViewer(StatusViewer): 

169 pass 

170 

171 

172class Task(object): 

173 def __init__( 

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

175 group=None): 

176 

177 self._id = id 

178 self._name = name 

179 self._condition = '' 

180 self._ispin = 0 

181 self._i = None 

182 self._n = n 

183 self._done = False 

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

185 self._state = state 

186 self._progress = progress 

187 self._logger = logger 

188 self._tcreate = time.time() 

189 self._group = group 

190 

191 def __enter__(self): 

192 return self 

193 

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

195 if type is None: 

196 self.done() 

197 else: 

198 self.fail() 

199 

200 def __call__(self, it): 

201 try: 

202 self._n = len(it) 

203 except TypeError: 

204 self._n = None 

205 

206 clean = False 

207 try: 

208 n = 0 

209 for obj in it: 

210 self.update(n) 

211 yield obj 

212 n += 1 

213 

214 self.update(n) 

215 clean = True 

216 

217 finally: 

218 if clean: 

219 self.done() 

220 else: 

221 self.fail() 

222 

223 def log(self, s): 

224 if self._logger is not None: 

225 self._logger.info(s) 

226 

227 def get_group_time_start(self): 

228 if self._group: 

229 return self._group.get_group_time_start() 

230 else: 

231 return self._tcreate 

232 

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

234 kwargs['group'] = self 

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

236 

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

238 self._state = 'working' 

239 

240 self._condition = condition 

241 

242 if i is not None: 

243 if self._n is not None: 

244 i = min(i, self._n) 

245 

246 self._i = i 

247 

248 self._progress._update() 

249 

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

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

252 

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

254 return 

255 

256 self._condition = condition 

257 self._state = 'done' 

258 self._progress._end(self) 

259 self.log(str(self)) 

260 

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

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

263 

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

265 return 

266 

267 self._condition = condition 

268 self._state = 'failed' 

269 self._progress._end(self) 

270 self.log(str(self)) 

271 

272 def _str_state(self): 

273 s = self._state 

274 if s == 'waiting': 

275 return ' ' 

276 elif s == 'working': 

277 self._ispin += 1 

278 return spinner[self._ispin % len(spinner)] + ' ' 

279 elif s == 'done': 

280 return symbol_done + ' ' 

281 elif s == 'failed': 

282 return symbol_failed + ' ' 

283 else: 

284 return '? ' 

285 

286 def _idisplay(self): 

287 i = self._i 

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

289 i = self._n 

290 return i 

291 

292 def _str_progress(self): 

293 if self._i is None: 

294 return self._state 

295 elif self._n is None: 

296 if self._state != 'working': 

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

298 else: 

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

300 else: 

301 if self._state == 'working': 

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

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

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

305 

306 elif self._state == 'failed': 

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

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

309 else: 

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

311 

312 def _str_percent(self): 

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

314 and self._i is not None: 

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

316 else: 

317 return '' 

318 

319 def _str_condition(self): 

320 if self._condition: 

321 return '%s' % self._condition 

322 else: 

323 return '' 

324 

325 def _str_bar(self): 

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

327 and self._i is not None: 

328 nb = 20 

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

330 ib = int(fb) 

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

332 if ib == 0 and ip == 0: 

333 ip = 1 # indication of start 

334 s = blocks[0] * ib 

335 if ib < nb: 

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

337 

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

339 return s 

340 else: 

341 return '' 

342 

343 def __str__(self): 

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

345 self._str_state(), 

346 self._name, 

347 self._str_progress(), 

348 self._str_bar(), 

349 self._str_percent(), 

350 self._str_condition()) 

351 

352 

353class Progress(object): 

354 

355 def __init__(self): 

356 self._current_id = 0 

357 self._current_group_id = 0 

358 self._tasks = {} 

359 self._tasks_done = [] 

360 self._last_update = 0.0 

361 self._terms = [] 

362 

363 def view(self, viewer=None): 

364 if g_force_viewer_off or self._terms: 

365 viewer = 'off' 

366 elif viewer is None: 

367 viewer = g_viewer 

368 

369 try: 

370 term = g_viewer_classes[viewer](self) 

371 except KeyError: 

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

373 

374 self._terms.append(term) 

375 return term 

376 

377 def hide(self, term): 

378 self._update(force=True) 

379 self._terms.remove(term) 

380 

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

382 self._current_id += 1 

383 task = Task( 

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

385 self._tasks[task._id] = task 

386 self._update(force=True) 

387 return task 

388 

389 def _end(self, task): 

390 del self._tasks[task._id] 

391 self._tasks_done.append(task) 

392 self._update(force=True) 

393 

394 def _update(self, force=False): 

395 now = time.time() 

396 if self._last_update + 0.1 < now or force: 

397 self._tasks_done = [] 

398 

399 lines = self._lines() 

400 for term in self._terms: 

401 term.draw(lines) 

402 

403 self._last_update = now 

404 

405 def _lines(self): 

406 task_ids = sorted(self._tasks) 

407 lines = [] 

408 for task_id in task_ids: 

409 task = self._tasks[task_id] 

410 lines.extend(str(task).splitlines()) 

411 

412 return lines 

413 

414 

415g_viewer_classes = { 

416 'terminal': TerminalStatusViewer, 

417 'log': LogStatusViewer, 

418 'off': DummyStatusViewer} 

419 

420progress = Progress()