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 distutils.util 

3import atexit 

4import io 

5import os 

6import re 

7import sys 

8import logging 

9import datetime 

10from python_utils.time import timedelta_to_seconds, epoch, format_time 

11from python_utils.converters import scale_1024 

12from python_utils.terminal import get_terminal_size 

13 

14import six 

15 

16 

17assert timedelta_to_seconds 

18assert get_terminal_size 

19assert format_time 

20assert scale_1024 

21assert epoch 

22 

23 

24ANSI_TERMS = ( 

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

26 '(sco)?ansi', 

27 'cygwin', 

28 'konsole', 

29 'linux', 

30 'rxvt', 

31 'screen', 

32 'tmux', 

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

34) 

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

36 

37 

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

39 if is_terminal is None: 

40 # Jupyter Notebooks define this variable and support progress bars 

41 if 'JPY_PARENT_PID' in os.environ: 

42 is_terminal = True 

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

44 # is no way to check. 

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

46 is_terminal = True 

47 

48 if is_terminal is None: 

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

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

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

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

53 # environment variables. 

54 try: 

55 is_tty = fd.isatty() 

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

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

58 is_terminal = True 

59 # ANSICON is a Windows ANSI compatible console 

60 elif 'ANSICON' in os.environ: 

61 is_terminal = True 

62 else: 

63 is_terminal = None 

64 except Exception: 

65 is_terminal = False 

66 

67 return is_terminal 

68 

69 

70def is_terminal(fd, is_terminal=None): 

71 if is_terminal is None: 

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

73 is_terminal = is_ansi_terminal(True) or None 

74 

75 if is_terminal is None: 

76 # Allow a environment variable override 

77 is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None) 

78 

79 if is_terminal is None: # pragma: no cover 

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

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

82 try: 

83 is_terminal = fd.isatty() 

84 except Exception: 

85 is_terminal = False 

86 

87 return is_terminal 

88 

89 

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

91 ''' 

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

93 

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

95 1.234 

96 >>> deltas_to_seconds(123) 

97 123.0 

98 >>> deltas_to_seconds(1.234) 

99 1.234 

100 >>> deltas_to_seconds(None, 1.234) 

101 1.234 

102 >>> deltas_to_seconds(0, 1.234) 

103 0.0 

104 >>> deltas_to_seconds() 

105 Traceback (most recent call last): 

106 ... 

107 ValueError: No valid deltas passed to `deltas_to_seconds` 

108 >>> deltas_to_seconds(None) 

109 Traceback (most recent call last): 

110 ... 

111 ValueError: No valid deltas passed to `deltas_to_seconds` 

112 >>> deltas_to_seconds(default=0.0) 

113 0.0 

114 ''' 

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

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

117 

118 for delta in deltas: 

119 if delta is None: 

120 continue 

121 if isinstance(delta, datetime.timedelta): 

122 return timedelta_to_seconds(delta) 

123 elif not isinstance(delta, float): 

124 return float(delta) 

125 else: 

126 return delta 

127 

128 if default is ValueError: 

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

130 else: 

131 return default 

132 

133 

134def no_color(value): 

135 ''' 

136 Return the `value` without ANSI escape codes 

137 

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

139 True 

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

141 'abc' 

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

143 'abc' 

144 ''' 

145 if isinstance(value, bytes): 

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

147 pattern = pattern.encode() 

148 replace = b'' 

149 assert isinstance(pattern, bytes) 

150 else: 

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

152 replace = '' 

153 

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

155 

156 

157def len_color(value): 

158 ''' 

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

160 

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

162 3 

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

164 3 

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

166 3 

167 ''' 

168 return len(no_color(value)) 

169 

170 

171def env_flag(name, default=None): 

172 ''' 

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

174 on/off, and returns it as a boolean 

175 

176 If the environt variable is not defined, or has an unknown value, returns 

177 `default` 

178 ''' 

179 try: 

180 return bool(distutils.util.strtobool(os.environ.get(name, ''))) 

181 except ValueError: 

182 return default 

183 

184 

185class WrappingIO: 

186 

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

188 self.buffer = six.StringIO() 

189 self.target = target 

190 self.capturing = capturing 

191 self.listeners = listeners 

192 self.needs_clear = False 

193 

194 def write(self, value): 

195 if self.capturing: 

196 self.buffer.write(value) 

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

198 self.needs_clear = True 

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

200 listener.update() 

201 else: 

202 self.target.write(value) 

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

204 self.flush_target() 

205 

206 def flush(self): 

207 self.buffer.flush() 

208 

209 def _flush(self): 

210 value = self.buffer.getvalue() 

211 if value: 

212 self.flush() 

213 self.target.write(value) 

214 self.buffer.seek(0) 

215 self.buffer.truncate(0) 

216 self.needs_clear = False 

217 

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

219 self.flush_target() 

220 

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

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

223 self.target.flush() 

224 

225 

226class StreamWrapper(object): 

227 '''Wrap stdout and stderr globally''' 

228 

229 def __init__(self): 

230 self.stdout = self.original_stdout = sys.stdout 

231 self.stderr = self.original_stderr = sys.stderr 

232 self.original_excepthook = sys.excepthook 

233 self.wrapped_stdout = 0 

234 self.wrapped_stderr = 0 

235 self.wrapped_excepthook = 0 

236 self.capturing = 0 

237 self.listeners = set() 

238 

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

240 self.wrap_stdout() 

241 

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

243 self.wrap_stderr() 

244 

245 def start_capturing(self, bar=None): 

246 if bar: # pragma: no branch 

247 self.listeners.add(bar) 

248 

249 self.capturing += 1 

250 self.update_capturing() 

251 

252 def stop_capturing(self, bar=None): 

253 if bar: # pragma: no branch 

254 try: 

255 self.listeners.remove(bar) 

256 except KeyError: 

257 pass 

258 

259 self.capturing -= 1 

260 self.update_capturing() 

261 

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

263 if isinstance(self.stdout, WrappingIO): 

264 self.stdout.capturing = self.capturing > 0 

265 

266 if isinstance(self.stderr, WrappingIO): 

267 self.stderr.capturing = self.capturing > 0 

268 

269 if self.capturing <= 0: 

270 self.flush() 

271 

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

273 if stdout: 

274 self.wrap_stdout() 

275 

276 if stderr: 

277 self.wrap_stderr() 

278 

279 def wrap_stdout(self): 

280 self.wrap_excepthook() 

281 

282 if not self.wrapped_stdout: 

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

284 listeners=self.listeners) 

285 self.wrapped_stdout += 1 

286 

287 return sys.stdout 

288 

289 def wrap_stderr(self): 

290 self.wrap_excepthook() 

291 

292 if not self.wrapped_stderr: 

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

294 listeners=self.listeners) 

295 self.wrapped_stderr += 1 

296 

297 return sys.stderr 

298 

299 def unwrap_excepthook(self): 

300 if self.wrapped_excepthook: 

301 self.wrapped_excepthook -= 1 

302 sys.excepthook = self.original_excepthook 

303 

304 def wrap_excepthook(self): 

305 if not self.wrapped_excepthook: 

306 logger.debug('wrapping excepthook') 

307 self.wrapped_excepthook += 1 

308 sys.excepthook = self.excepthook 

309 

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

311 if stdout: 

312 self.unwrap_stdout() 

313 

314 if stderr: 

315 self.unwrap_stderr() 

316 

317 def unwrap_stdout(self): 

318 if self.wrapped_stdout > 1: 

319 self.wrapped_stdout -= 1 

320 else: 

321 sys.stdout = self.original_stdout 

322 self.wrapped_stdout = 0 

323 

324 def unwrap_stderr(self): 

325 if self.wrapped_stderr > 1: 

326 self.wrapped_stderr -= 1 

327 else: 

328 sys.stderr = self.original_stderr 

329 self.wrapped_stderr = 0 

330 

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

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

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

334 return stderr_needs_clear or stdout_needs_clear 

335 

336 def flush(self): 

337 if self.wrapped_stdout: # pragma: no branch 

338 try: 

339 self.stdout._flush() 

340 except (io.UnsupportedOperation, 

341 AttributeError): # pragma: no cover 

342 self.wrapped_stdout = False 

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

344 sys.stdout) 

345 

346 if self.wrapped_stderr: # pragma: no branch 

347 try: 

348 self.stderr._flush() 

349 except (io.UnsupportedOperation, 

350 AttributeError): # pragma: no cover 

351 self.wrapped_stderr = False 

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

353 sys.stderr) 

354 

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

356 self.original_excepthook(exc_type, exc_value, exc_traceback) 

357 self.flush() 

358 

359 

360class AttributeDict(dict): 

361 ''' 

362 A dict that can be accessed with .attribute 

363 

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

365 

366 # Reading 

367 >>> attrs['spam'] 

368 123 

369 >>> attrs.spam 

370 123 

371 

372 # Read after update using attribute 

373 >>> attrs.spam = 456 

374 >>> attrs['spam'] 

375 456 

376 >>> attrs.spam 

377 456 

378 

379 # Read after update using dict access 

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

381 >>> attrs['spam'] 

382 123 

383 >>> attrs.spam 

384 123 

385 

386 # Read after update using dict access 

387 >>> del attrs.spam 

388 >>> attrs['spam'] 

389 Traceback (most recent call last): 

390 ... 

391 KeyError: 'spam' 

392 >>> attrs.spam 

393 Traceback (most recent call last): 

394 ... 

395 AttributeError: No such attribute: spam 

396 >>> del attrs.spam 

397 Traceback (most recent call last): 

398 ... 

399 AttributeError: No such attribute: spam 

400 ''' 

401 def __getattr__(self, name): 

402 if name in self: 

403 return self[name] 

404 else: 

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

406 

407 def __setattr__(self, name, value): 

408 self[name] = value 

409 

410 def __delattr__(self, name): 

411 if name in self: 

412 del self[name] 

413 else: 

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

415 

416 

417logger = logging.getLogger(__name__) 

418streams = StreamWrapper() 

419atexit.register(streams.flush)