Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from __future__ import absolute_import 

2import atexit 

3import io 

4import os 

5import re 

6import sys 

7import logging 

8import datetime 

9from python_utils.time import timedelta_to_seconds, epoch, format_time 

10from python_utils.converters import scale_1024 

11from python_utils.terminal import get_terminal_size 

12 

13import six 

14 

15 

16assert timedelta_to_seconds 

17assert get_terminal_size 

18assert format_time 

19assert scale_1024 

20assert epoch 

21 

22 

23ANSI_TERMS = ( 

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

25 '(sco)?ansi', 

26 'cygwin', 

27 'konsole', 

28 'linux', 

29 'rxvt', 

30 'screen', 

31 'tmux', 

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

33) 

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

35 

36 

37def is_ansi_terminal(fd, is_terminal=None): # pragma: no cover 

38 if is_terminal is None: 

39 # Jupyter Notebooks define this variable and support progress bars 

40 if 'JPY_PARENT_PID' in os.environ: 

41 is_terminal = True 

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

43 # is no way to check. 

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

45 is_terminal = True 

46 

47 if is_terminal is None: 

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

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

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

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

52 # environment variables. 

53 try: 

54 is_tty = fd.isatty() 

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

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

57 is_terminal = True 

58 # ANSICON is a Windows ANSI compatible console 

59 elif 'ANSICON' in os.environ: 

60 is_terminal = True 

61 else: 

62 is_terminal = None 

63 except Exception: 

64 is_terminal = False 

65 

66 return is_terminal 

67 

68 

69def is_terminal(fd, is_terminal=None): 

70 if is_terminal is None: 

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

72 is_terminal = is_ansi_terminal(True) or None 

73 

74 if is_terminal is None: 

75 # Allow a environment variable override 

76 is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None) 

77 

78 if is_terminal is None: # pragma: no cover 

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

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

81 try: 

82 is_terminal = fd.isatty() 

83 except Exception: 

84 is_terminal = False 

85 

86 return is_terminal 

87 

88 

89def deltas_to_seconds(*deltas, **kwargs): # default=ValueError): 

90 ''' 

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

92 

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

94 1.234 

95 >>> deltas_to_seconds(123) 

96 123.0 

97 >>> deltas_to_seconds(1.234) 

98 1.234 

99 >>> deltas_to_seconds(None, 1.234) 

100 1.234 

101 >>> deltas_to_seconds(0, 1.234) 

102 0.0 

103 >>> deltas_to_seconds() 

104 Traceback (most recent call last): 

105 ... 

106 ValueError: No valid deltas passed to `deltas_to_seconds` 

107 >>> deltas_to_seconds(None) 

108 Traceback (most recent call last): 

109 ... 

110 ValueError: No valid deltas passed to `deltas_to_seconds` 

111 >>> deltas_to_seconds(default=0.0) 

112 0.0 

113 ''' 

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

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

116 

117 for delta in deltas: 

118 if delta is None: 

119 continue 

120 if isinstance(delta, datetime.timedelta): 

121 return timedelta_to_seconds(delta) 

122 elif not isinstance(delta, float): 

123 return float(delta) 

124 else: 

125 return delta 

126 

127 if default is ValueError: 

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

129 else: 

130 return default 

131 

132 

133def no_color(value): 

134 ''' 

135 Return the `value` without ANSI escape codes 

136 

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

138 True 

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

140 'abc' 

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

142 'abc' 

143 ''' 

144 if isinstance(value, bytes): 

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

146 pattern = pattern.encode() 

147 replace = b'' 

148 assert isinstance(pattern, bytes) 

149 else: 

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

151 replace = '' 

152 

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

154 

155 

156def len_color(value): 

157 ''' 

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

159 

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

161 3 

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

163 3 

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

165 3 

166 ''' 

167 return len(no_color(value)) 

168 

169 

170def env_flag(name, default=None): 

171 ''' 

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

173 on/off, and returns it as a boolean 

174 

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

176 returns `default` 

177 ''' 

178 v = os.getenv(name) 

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

180 return True 

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

182 return False 

183 return default 

184 

185 

186class WrappingIO: 

187 

188 def __init__(self, target, capturing=False, listeners=set()): 

189 self.buffer = six.StringIO() 

190 self.target = target 

191 self.capturing = capturing 

192 self.listeners = listeners 

193 self.needs_clear = False 

194 

195 def isatty(self): # pragma: no cover 

196 return self.target.isatty() 

197 

198 def write(self, value): 

199 if self.capturing: 

200 self.buffer.write(value) 

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

202 self.needs_clear = True 

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

204 listener.update() 

205 else: 

206 self.target.write(value) 

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

208 self.flush_target() 

209 

210 def flush(self): 

211 self.buffer.flush() 

212 

213 def _flush(self): 

214 value = self.buffer.getvalue() 

215 if value: 

216 self.flush() 

217 self.target.write(value) 

218 self.buffer.seek(0) 

219 self.buffer.truncate(0) 

220 self.needs_clear = False 

221 

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

223 self.flush_target() 

224 

225 def flush_target(self): # pragma: no cover 

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

227 self.target.flush() 

228 

229 

230class StreamWrapper(object): 

231 '''Wrap stdout and stderr globally''' 

232 

233 def __init__(self): 

234 self.stdout = self.original_stdout = sys.stdout 

235 self.stderr = self.original_stderr = sys.stderr 

236 self.original_excepthook = sys.excepthook 

237 self.wrapped_stdout = 0 

238 self.wrapped_stderr = 0 

239 self.wrapped_excepthook = 0 

240 self.capturing = 0 

241 self.listeners = set() 

242 

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

244 self.wrap_stdout() 

245 

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

247 self.wrap_stderr() 

248 

249 def start_capturing(self, bar=None): 

250 if bar: # pragma: no branch 

251 self.listeners.add(bar) 

252 

253 self.capturing += 1 

254 self.update_capturing() 

255 

256 def stop_capturing(self, bar=None): 

257 if bar: # pragma: no branch 

258 try: 

259 self.listeners.remove(bar) 

260 except KeyError: 

261 pass 

262 

263 self.capturing -= 1 

264 self.update_capturing() 

265 

266 def update_capturing(self): # pragma: no cover 

267 if isinstance(self.stdout, WrappingIO): 

268 self.stdout.capturing = self.capturing > 0 

269 

270 if isinstance(self.stderr, WrappingIO): 

271 self.stderr.capturing = self.capturing > 0 

272 

273 if self.capturing <= 0: 

274 self.flush() 

275 

276 def wrap(self, stdout=False, stderr=False): 

277 if stdout: 

278 self.wrap_stdout() 

279 

280 if stderr: 

281 self.wrap_stderr() 

282 

283 def wrap_stdout(self): 

284 self.wrap_excepthook() 

285 

286 if not self.wrapped_stdout: 

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

288 listeners=self.listeners) 

289 self.wrapped_stdout += 1 

290 

291 return sys.stdout 

292 

293 def wrap_stderr(self): 

294 self.wrap_excepthook() 

295 

296 if not self.wrapped_stderr: 

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

298 listeners=self.listeners) 

299 self.wrapped_stderr += 1 

300 

301 return sys.stderr 

302 

303 def unwrap_excepthook(self): 

304 if self.wrapped_excepthook: 

305 self.wrapped_excepthook -= 1 

306 sys.excepthook = self.original_excepthook 

307 

308 def wrap_excepthook(self): 

309 if not self.wrapped_excepthook: 

310 logger.debug('wrapping excepthook') 

311 self.wrapped_excepthook += 1 

312 sys.excepthook = self.excepthook 

313 

314 def unwrap(self, stdout=False, stderr=False): 

315 if stdout: 

316 self.unwrap_stdout() 

317 

318 if stderr: 

319 self.unwrap_stderr() 

320 

321 def unwrap_stdout(self): 

322 if self.wrapped_stdout > 1: 

323 self.wrapped_stdout -= 1 

324 else: 

325 sys.stdout = self.original_stdout 

326 self.wrapped_stdout = 0 

327 

328 def unwrap_stderr(self): 

329 if self.wrapped_stderr > 1: 

330 self.wrapped_stderr -= 1 

331 else: 

332 sys.stderr = self.original_stderr 

333 self.wrapped_stderr = 0 

334 

335 def needs_clear(self): # pragma: no cover 

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

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

338 return stderr_needs_clear or stdout_needs_clear 

339 

340 def flush(self): 

341 if self.wrapped_stdout: # pragma: no branch 

342 try: 

343 self.stdout._flush() 

344 except (io.UnsupportedOperation, 

345 AttributeError): # pragma: no cover 

346 self.wrapped_stdout = False 

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

348 sys.stdout) 

349 

350 if self.wrapped_stderr: # pragma: no branch 

351 try: 

352 self.stderr._flush() 

353 except (io.UnsupportedOperation, 

354 AttributeError): # pragma: no cover 

355 self.wrapped_stderr = False 

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

357 sys.stderr) 

358 

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

360 self.original_excepthook(exc_type, exc_value, exc_traceback) 

361 self.flush() 

362 

363 

364class AttributeDict(dict): 

365 ''' 

366 A dict that can be accessed with .attribute 

367 

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

369 

370 # Reading 

371 >>> attrs['spam'] 

372 123 

373 >>> attrs.spam 

374 123 

375 

376 # Read after update using attribute 

377 >>> attrs.spam = 456 

378 >>> attrs['spam'] 

379 456 

380 >>> attrs.spam 

381 456 

382 

383 # Read after update using dict access 

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

385 >>> attrs['spam'] 

386 123 

387 >>> attrs.spam 

388 123 

389 

390 # Read after update using dict access 

391 >>> del attrs.spam 

392 >>> attrs['spam'] 

393 Traceback (most recent call last): 

394 ... 

395 KeyError: 'spam' 

396 >>> attrs.spam 

397 Traceback (most recent call last): 

398 ... 

399 AttributeError: No such attribute: spam 

400 >>> del attrs.spam 

401 Traceback (most recent call last): 

402 ... 

403 AttributeError: No such attribute: spam 

404 ''' 

405 def __getattr__(self, name): 

406 if name in self: 

407 return self[name] 

408 else: 

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

410 

411 def __setattr__(self, name, value): 

412 self[name] = value 

413 

414 def __delattr__(self, name): 

415 if name in self: 

416 del self[name] 

417 else: 

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

419 

420 

421logger = logging.getLogger(__name__) 

422streams = StreamWrapper() 

423atexit.register(streams.flush)