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 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 if 'JPY_PARENT_PID' in os.environ: 

40 is_terminal = True 

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

42 is_terminal = True 

43 

44 if is_terminal is None: 

45 try: 

46 is_tty = fd.isatty() 

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

48 is_terminal = True 

49 elif 'ANSICON' in os.environ: 

50 is_terminal = True 

51 else: 

52 is_terminal = False 

53 except Exception: 

54 is_terminal = False 

55 

56 return is_terminal 

57 

58 

59def is_terminal(fd, is_terminal=None): 

60 if is_terminal is None: 

61 is_terminal = is_ansi_terminal(True) or None 

62 

63 if is_terminal is None: 

64 is_terminal = env_flag('PROGRESSBAR_IS_TERMINAL', None) 

65 

66 if is_terminal is None: # pragma: no cover 

67 try: 

68 is_terminal = fd.isatty() 

69 except Exception: 

70 is_terminal = False 

71 

72 return is_terminal 

73 

74 

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

76 ''' 

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

78 

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

80 1.234 

81 >>> deltas_to_seconds(123) 

82 123.0 

83 >>> deltas_to_seconds(1.234) 

84 1.234 

85 >>> deltas_to_seconds(None, 1.234) 

86 1.234 

87 >>> deltas_to_seconds(0, 1.234) 

88 0.0 

89 >>> deltas_to_seconds() 

90 Traceback (most recent call last): 

91 ... 

92 ValueError: No valid deltas passed to `deltas_to_seconds` 

93 >>> deltas_to_seconds(None) 

94 Traceback (most recent call last): 

95 ... 

96 ValueError: No valid deltas passed to `deltas_to_seconds` 

97 >>> deltas_to_seconds(default=0.0) 

98 0.0 

99 ''' 

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

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

102 

103 for delta in deltas: 

104 if delta is None: 

105 continue 

106 if isinstance(delta, datetime.timedelta): 

107 return timedelta_to_seconds(delta) 

108 elif not isinstance(delta, float): 

109 return float(delta) 

110 else: 

111 return delta 

112 

113 if default is ValueError: 

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

115 else: 

116 return default 

117 

118 

119def no_color(value): 

120 ''' 

121 Return the `value` without ANSI escape codes 

122 

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

124 True 

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

126 'abc' 

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

128 'abc' 

129 ''' 

130 if isinstance(value, bytes): 

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

132 pattern = pattern.encode() 

133 replace = b'' 

134 assert isinstance(pattern, bytes) 

135 else: 

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

137 replace = '' 

138 

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

140 

141 

142def len_color(value): 

143 ''' 

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

145 

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

147 3 

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

149 3 

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

151 3 

152 ''' 

153 return len(no_color(value)) 

154 

155 

156def env_flag(name, default=None): 

157 ''' 

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

159 on/off, and returns it as a boolean 

160 

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

162 `default` 

163 ''' 

164 try: 

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

166 except ValueError: 

167 return default 

168 

169 

170class WrappingIO: 

171 

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

173 self.buffer = six.StringIO() 

174 self.target = target 

175 self.capturing = capturing 

176 self.listeners = listeners 

177 self.needs_clear = False 

178 

179 def write(self, value): 

180 if self.capturing: 

181 self.buffer.write(value) 

182 if '\n' in value: 

183 self.needs_clear = True 

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

185 listener.update() 

186 else: 

187 self.target.write(value) 

188 

189 def flush(self): 

190 self.needs_clear = False 

191 self.buffer.flush() 

192 

193 def _flush(self): 

194 value = self.buffer.getvalue() 

195 if value: 

196 self.flush() 

197 self.target.write(value) 

198 self.buffer.seek(0) 

199 self.buffer.truncate(0) 

200 

201 

202class StreamWrapper(object): 

203 '''Wrap stdout and stderr globally''' 

204 

205 def __init__(self): 

206 self.stdout = self.original_stdout = sys.stdout 

207 self.stderr = self.original_stderr = sys.stderr 

208 self.original_excepthook = sys.excepthook 

209 self.wrapped_stdout = 0 

210 self.wrapped_stderr = 0 

211 self.wrapped_excepthook = 0 

212 self.capturing = 0 

213 self.listeners = set() 

214 

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

216 self.wrap_stdout() 

217 

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

219 self.wrap_stderr() 

220 

221 def start_capturing(self, bar=None): 

222 if bar: # pragma: no branch 

223 self.listeners.add(bar) 

224 

225 self.capturing += 1 

226 self.update_capturing() 

227 

228 def stop_capturing(self, bar=None): 

229 if bar: # pragma: no branch 

230 try: 

231 self.listeners.remove(bar) 

232 except KeyError: 

233 pass 

234 

235 self.capturing -= 1 

236 self.update_capturing() 

237 

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

239 if isinstance(self.stdout, WrappingIO): 

240 self.stdout.capturing = self.capturing > 0 

241 

242 if isinstance(self.stderr, WrappingIO): 

243 self.stderr.capturing = self.capturing > 0 

244 

245 if self.capturing <= 0: 

246 self.flush() 

247 

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

249 if stdout: 

250 self.wrap_stdout() 

251 

252 if stderr: 

253 self.wrap_stderr() 

254 

255 def wrap_stdout(self): 

256 self.wrap_excepthook() 

257 

258 if not self.wrapped_stdout: 

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

260 listeners=self.listeners) 

261 self.wrapped_stdout += 1 

262 

263 return sys.stdout 

264 

265 def wrap_stderr(self): 

266 self.wrap_excepthook() 

267 

268 if not self.wrapped_stderr: 

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

270 listeners=self.listeners) 

271 self.wrapped_stderr += 1 

272 

273 return sys.stderr 

274 

275 def unwrap_excepthook(self): 

276 if self.wrapped_excepthook: 

277 self.wrapped_excepthook -= 1 

278 sys.excepthook = self.original_excepthook 

279 

280 def wrap_excepthook(self): 

281 if not self.wrapped_excepthook: 

282 logger.debug('wrapping excepthook') 

283 self.wrapped_excepthook += 1 

284 sys.excepthook = self.excepthook 

285 

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

287 if stdout: 

288 self.unwrap_stdout() 

289 

290 if stderr: 

291 self.unwrap_stderr() 

292 

293 def unwrap_stdout(self): 

294 if self.wrapped_stdout > 1: 

295 self.wrapped_stdout -= 1 

296 else: 

297 sys.stdout = self.original_stdout 

298 self.wrapped_stdout = 0 

299 

300 def unwrap_stderr(self): 

301 if self.wrapped_stderr > 1: 

302 self.wrapped_stderr -= 1 

303 else: 

304 sys.stderr = self.original_stderr 

305 self.wrapped_stderr = 0 

306 

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

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

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

310 return stderr_needs_clear or stdout_needs_clear 

311 

312 def flush(self): 

313 if self.wrapped_stdout: # pragma: no branch 

314 try: 

315 self.stdout._flush() 

316 except (io.UnsupportedOperation, 

317 AttributeError): # pragma: no cover 

318 self.wrapped_stdout = False 

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

320 sys.stdout) 

321 

322 if self.wrapped_stderr: # pragma: no branch 

323 try: 

324 self.stderr._flush() 

325 except (io.UnsupportedOperation, 

326 AttributeError): # pragma: no cover 

327 self.wrapped_stderr = False 

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

329 sys.stderr) 

330 

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

332 self.original_excepthook(exc_type, exc_value, exc_traceback) 

333 self.flush() 

334 

335 

336class AttributeDict(dict): 

337 ''' 

338 A dict that can be accessed with .attribute 

339 

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

341 

342 # Reading 

343 >>> attrs['spam'] 

344 123 

345 >>> attrs.spam 

346 123 

347 

348 # Read after update using attribute 

349 >>> attrs.spam = 456 

350 >>> attrs['spam'] 

351 456 

352 >>> attrs.spam 

353 456 

354 

355 # Read after update using dict access 

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

357 >>> attrs['spam'] 

358 123 

359 >>> attrs.spam 

360 123 

361 

362 # Read after update using dict access 

363 >>> del attrs.spam 

364 >>> attrs['spam'] 

365 Traceback (most recent call last): 

366 ... 

367 KeyError: 'spam' 

368 >>> attrs.spam 

369 Traceback (most recent call last): 

370 ... 

371 AttributeError: No such attribute: spam 

372 >>> del attrs.spam 

373 Traceback (most recent call last): 

374 ... 

375 AttributeError: No such attribute: spam 

376 ''' 

377 def __getattr__(self, name): 

378 if name in self: 

379 return self[name] 

380 else: 

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

382 

383 def __setattr__(self, name, value): 

384 self[name] = value 

385 

386 def __delattr__(self, name): 

387 if name in self: 

388 del self[name] 

389 else: 

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

391 

392 

393logger = logging.getLogger(__name__) 

394streams = StreamWrapper()