Coverage for /Volumes/workspace/python-progressbar/.tox/py37/lib/python3.7/site-packages/progressbar/utils.py: 93%

177 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-26 14:16 +0200

1from __future__ import annotations 

2 

3import atexit 

4import datetime 

5import io 

6import logging 

7import os 

8import re 

9import sys 

10 

11from python_utils import types 

12from python_utils.converters import scale_1024 

13from python_utils.terminal import get_terminal_size 

14from python_utils.time import epoch 

15from python_utils.time import format_time 

16from python_utils.time import timedelta_to_seconds 

17 

18if types.TYPE_CHECKING: 

19 from .bar import ProgressBar 

20 

21assert timedelta_to_seconds 

22assert get_terminal_size 

23assert format_time 

24assert scale_1024 

25assert epoch 

26 

27ANSI_TERMS = ( 

28 '([xe]|bv)term', 

29 '(sco)?ansi', 

30 'cygwin', 

31 'konsole', 

32 'linux', 

33 'rxvt', 

34 'screen', 

35 'tmux', 

36 'vt(10[02]|220|320)', 

37) 

38ANSI_TERM_RE = re.compile('^({})'.format('|'.join(ANSI_TERMS)), re.IGNORECASE) 

39 

40 

41def is_ansi_terminal(fd: types.IO, is_terminal: bool | None = None) \ 

42 -> bool: # pragma: no cover 

43 if is_terminal is None: 

44 # Jupyter Notebooks define this variable and support progress bars 

45 if 'JPY_PARENT_PID' in os.environ: 

46 is_terminal = True 

47 # This works for newer versions of pycharm only. older versions there 

48 # is no way to check. 

49 elif os.environ.get('PYCHARM_HOSTED') == '1': 

50 is_terminal = True 

51 

52 if is_terminal is None: 

53 # check if we are writing to a terminal or not. typically a file object 

54 # is going to return False if the instance has been overridden and 

55 # isatty has not been defined we have no way of knowing so we will not 

56 # use ansi. ansi terminals will typically define one of the 2 

57 # environment variables. 

58 try: 

59 is_tty = fd.isatty() 

60 # Try and match any of the huge amount of Linux/Unix ANSI consoles 

61 if is_tty and ANSI_TERM_RE.match(os.environ.get('TERM', '')): 

62 is_terminal = True 

63 # ANSICON is a Windows ANSI compatible console 

64 elif 'ANSICON' in os.environ: 

65 is_terminal = True 

66 else: 

67 is_terminal = None 

68 except Exception: 

69 is_terminal = False 

70 

71 return is_terminal 

72 

73 

74def is_terminal(fd: types.IO, is_terminal: bool | None = None) -> bool: 

75 if is_terminal is None: 

76 # Full ansi support encompasses what we expect from a terminal 

77 is_terminal = is_ansi_terminal(True) or None 

78 

79 if is_terminal is None: 

80 # Allow a environment variable override 

81 is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None) 

82 

83 if is_terminal is None: # pragma: no cover 

84 # Bare except because a lot can go wrong on different systems. If we do 

85 # get a TTY we know this is a valid terminal 

86 try: 

87 is_terminal = fd.isatty() 

88 except Exception: 

89 is_terminal = False 

90 

91 return is_terminal 

92 

93 

94def deltas_to_seconds(*deltas, 

95 **kwargs) -> int | float | None: # default=ValueError): 

96 ''' 

97 Convert timedeltas and seconds as int to seconds as float while coalescing 

98 

99 >>> deltas_to_seconds(datetime.timedelta(seconds=1, milliseconds=234)) 

100 1.234 

101 >>> deltas_to_seconds(123) 

102 123.0 

103 >>> deltas_to_seconds(1.234) 

104 1.234 

105 >>> deltas_to_seconds(None, 1.234) 

106 1.234 

107 >>> deltas_to_seconds(0, 1.234) 

108 0.0 

109 >>> deltas_to_seconds() 

110 Traceback (most recent call last): 

111 ... 

112 ValueError: No valid deltas passed to `deltas_to_seconds` 

113 >>> deltas_to_seconds(None) 

114 Traceback (most recent call last): 

115 ... 

116 ValueError: No valid deltas passed to `deltas_to_seconds` 

117 >>> deltas_to_seconds(default=0.0) 

118 0.0 

119 ''' 

120 default = kwargs.pop('default', ValueError) 

121 assert not kwargs, 'Only the `default` keyword argument is supported' 

122 

123 for delta in deltas: 

124 if delta is None: 

125 continue 

126 if isinstance(delta, datetime.timedelta): 

127 return timedelta_to_seconds(delta) 

128 elif not isinstance(delta, float): 

129 return float(delta) 

130 else: 

131 return delta 

132 

133 if default is ValueError: 

134 raise ValueError('No valid deltas passed to `deltas_to_seconds`') 

135 else: 

136 return default 

137 

138 

139def no_color(value: types.StringTypes) -> types.StringTypes: 

140 ''' 

141 Return the `value` without ANSI escape codes 

142 

143 >>> no_color(b'\u001b[1234]abc') == b'abc' 

144 True 

145 >>> str(no_color(u'\u001b[1234]abc')) 

146 'abc' 

147 >>> str(no_color('\u001b[1234]abc')) 

148 'abc' 

149 ''' 

150 if isinstance(value, bytes): 

151 pattern = '\\\u001b\\[.*?[@-~]' 

152 pattern = pattern.encode() 

153 replace = b'' 

154 assert isinstance(pattern, bytes) 

155 else: 

156 pattern = u'\x1b\\[.*?[@-~]' 

157 replace = '' 

158 

159 return re.sub(pattern, replace, value) 

160 

161 

162def len_color(value: types.StringTypes) -> int: 

163 ''' 

164 Return the length of `value` without ANSI escape codes 

165 

166 >>> len_color(b'\u001b[1234]abc') 

167 3 

168 >>> len_color(u'\u001b[1234]abc') 

169 3 

170 >>> len_color('\u001b[1234]abc') 

171 3 

172 ''' 

173 return len(no_color(value)) 

174 

175 

176def env_flag(name: str, default: bool | None = None) -> bool | None: 

177 ''' 

178 Accepts environt variables formatted as y/n, yes/no, 1/0, true/false, 

179 on/off, and returns it as a boolean 

180 

181 If the environment variable is not defined, or has an unknown value, 

182 returns `default` 

183 ''' 

184 v = os.getenv(name) 

185 if v and v.lower() in ('y', 'yes', 't', 'true', 'on', '1'): 

186 return True 

187 if v and v.lower() in ('n', 'no', 'f', 'false', 'off', '0'): 

188 return False 

189 return default 

190 

191 

192class WrappingIO: 

193 

194 def __init__(self, target: types.IO, capturing: bool = False, 

195 listeners: types.Set[ProgressBar] = None) -> None: 

196 self.buffer = io.StringIO() 

197 self.target = target 

198 self.capturing = capturing 

199 self.listeners = listeners or set() 

200 self.needs_clear = False 

201 

202 def __getattr__(self, name): # pragma: no cover 

203 return getattr(self.target, name) 

204 

205 def write(self, value: str) -> None: 

206 if self.capturing: 

207 self.buffer.write(value) 

208 if '\n' in value: # pragma: no branch 

209 self.needs_clear = True 

210 for listener in self.listeners: # pragma: no branch 

211 listener.update() 

212 else: 

213 self.target.write(value) 

214 if '\n' in value: # pragma: no branch 

215 self.flush_target() 

216 

217 def flush(self) -> None: 

218 self.buffer.flush() 

219 

220 def _flush(self) -> None: 

221 value = self.buffer.getvalue() 

222 if value: 

223 self.flush() 

224 self.target.write(value) 

225 self.buffer.seek(0) 

226 self.buffer.truncate(0) 

227 self.needs_clear = False 

228 

229 # when explicitly flushing, always flush the target as well 

230 self.flush_target() 

231 

232 def flush_target(self) -> None: # pragma: no cover 

233 if not self.target.closed and getattr(self.target, 'flush'): 

234 self.target.flush() 

235 

236 

237class StreamWrapper: 

238 '''Wrap stdout and stderr globally''' 

239 

240 def __init__(self): 

241 self.stdout = self.original_stdout = sys.stdout 

242 self.stderr = self.original_stderr = sys.stderr 

243 self.original_excepthook = sys.excepthook 

244 self.wrapped_stdout = 0 

245 self.wrapped_stderr = 0 

246 self.wrapped_excepthook = 0 

247 self.capturing = 0 

248 self.listeners = set() 

249 

250 if env_flag('WRAP_STDOUT', default=False): # pragma: no cover 

251 self.wrap_stdout() 

252 

253 if env_flag('WRAP_STDERR', default=False): # pragma: no cover 

254 self.wrap_stderr() 

255 

256 def start_capturing(self, bar: ProgressBar | None = None) -> None: 

257 if bar: # pragma: no branch 

258 self.listeners.add(bar) 

259 

260 self.capturing += 1 

261 self.update_capturing() 

262 

263 def stop_capturing(self, bar: ProgressBar | None = None) -> None: 

264 if bar: # pragma: no branch 

265 try: 

266 self.listeners.remove(bar) 

267 except KeyError: 

268 pass 

269 

270 self.capturing -= 1 

271 self.update_capturing() 

272 

273 def update_capturing(self) -> None: # pragma: no cover 

274 if isinstance(self.stdout, WrappingIO): 

275 self.stdout.capturing = self.capturing > 0 

276 

277 if isinstance(self.stderr, WrappingIO): 

278 self.stderr.capturing = self.capturing > 0 

279 

280 if self.capturing <= 0: 

281 self.flush() 

282 

283 def wrap(self, stdout: bool = False, stderr: bool = False) -> None: 

284 if stdout: 

285 self.wrap_stdout() 

286 

287 if stderr: 

288 self.wrap_stderr() 

289 

290 def wrap_stdout(self) -> types.IO: 

291 self.wrap_excepthook() 

292 

293 if not self.wrapped_stdout: 

294 self.stdout = sys.stdout = WrappingIO(self.original_stdout, 

295 listeners=self.listeners) 

296 self.wrapped_stdout += 1 

297 

298 return sys.stdout 

299 

300 def wrap_stderr(self) -> types.IO: 

301 self.wrap_excepthook() 

302 

303 if not self.wrapped_stderr: 

304 self.stderr = sys.stderr = WrappingIO(self.original_stderr, 

305 listeners=self.listeners) 

306 self.wrapped_stderr += 1 

307 

308 return sys.stderr 

309 

310 def unwrap_excepthook(self) -> None: 

311 if self.wrapped_excepthook: 

312 self.wrapped_excepthook -= 1 

313 sys.excepthook = self.original_excepthook 

314 

315 def wrap_excepthook(self) -> None: 

316 if not self.wrapped_excepthook: 

317 logger.debug('wrapping excepthook') 

318 self.wrapped_excepthook += 1 

319 sys.excepthook = self.excepthook 

320 

321 def unwrap(self, stdout: bool = False, stderr: bool = False) -> None: 

322 if stdout: 

323 self.unwrap_stdout() 

324 

325 if stderr: 

326 self.unwrap_stderr() 

327 

328 def unwrap_stdout(self) -> None: 

329 if self.wrapped_stdout > 1: 

330 self.wrapped_stdout -= 1 

331 else: 

332 sys.stdout = self.original_stdout 

333 self.wrapped_stdout = 0 

334 

335 def unwrap_stderr(self) -> None: 

336 if self.wrapped_stderr > 1: 

337 self.wrapped_stderr -= 1 

338 else: 

339 sys.stderr = self.original_stderr 

340 self.wrapped_stderr = 0 

341 

342 def needs_clear(self) -> bool: # pragma: no cover 

343 stdout_needs_clear = getattr(self.stdout, 'needs_clear', False) 

344 stderr_needs_clear = getattr(self.stderr, 'needs_clear', False) 

345 return stderr_needs_clear or stdout_needs_clear 

346 

347 def flush(self) -> None: 

348 if self.wrapped_stdout: # pragma: no branch 

349 try: 

350 self.stdout._flush() 

351 except (io.UnsupportedOperation, 

352 AttributeError): # pragma: no cover 

353 self.wrapped_stdout = False 

354 logger.warn('Disabling stdout redirection, %r is not seekable', 

355 sys.stdout) 

356 

357 if self.wrapped_stderr: # pragma: no branch 

358 try: 

359 self.stderr._flush() 

360 except (io.UnsupportedOperation, 

361 AttributeError): # pragma: no cover 

362 self.wrapped_stderr = False 

363 logger.warn('Disabling stderr redirection, %r is not seekable', 

364 sys.stderr) 

365 

366 def excepthook(self, exc_type, exc_value, exc_traceback): 

367 self.original_excepthook(exc_type, exc_value, exc_traceback) 

368 self.flush() 

369 

370 

371class AttributeDict(dict): 

372 ''' 

373 A dict that can be accessed with .attribute 

374 

375 >>> attrs = AttributeDict(spam=123) 

376 

377 # Reading 

378 

379 >>> attrs['spam'] 

380 123 

381 >>> attrs.spam 

382 123 

383 

384 # Read after update using attribute 

385 

386 >>> attrs.spam = 456 

387 >>> attrs['spam'] 

388 456 

389 >>> attrs.spam 

390 456 

391 

392 # Read after update using dict access 

393 

394 >>> attrs['spam'] = 123 

395 >>> attrs['spam'] 

396 123 

397 >>> attrs.spam 

398 123 

399 

400 # Read after update using dict access 

401 

402 >>> del attrs.spam 

403 >>> attrs['spam'] 

404 Traceback (most recent call last): 

405 ... 

406 KeyError: 'spam' 

407 >>> attrs.spam 

408 Traceback (most recent call last): 

409 ... 

410 AttributeError: No such attribute: spam 

411 >>> del attrs.spam 

412 Traceback (most recent call last): 

413 ... 

414 AttributeError: No such attribute: spam 

415 ''' 

416 

417 def __getattr__(self, name: str) -> int: 

418 if name in self: 

419 return self[name] 

420 else: 

421 raise AttributeError("No such attribute: " + name) 

422 

423 def __setattr__(self, name: str, value: int) -> None: 

424 self[name] = value 

425 

426 def __delattr__(self, name: str) -> None: 

427 if name in self: 

428 del self[name] 

429 else: 

430 raise AttributeError("No such attribute: " + name) 

431 

432 

433logger = logging.getLogger(__name__) 

434streams = StreamWrapper() 

435atexit.register(streams.flush)