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 # 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 environt variable is not defined, or has an unknown value, returns 

176 `default` 

177 ''' 

178 try: 

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

180 except ValueError: 

181 return default 

182 

183 

184class WrappingIO: 

185 

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

187 self.buffer = six.StringIO() 

188 self.target = target 

189 self.capturing = capturing 

190 self.listeners = listeners 

191 self.needs_clear = False 

192 

193 def write(self, value): 

194 if self.capturing: 

195 self.buffer.write(value) 

196 if '\n' in value: 

197 self.needs_clear = True 

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

199 listener.update() 

200 else: 

201 self.target.write(value) 

202 

203 def flush(self): 

204 self.buffer.flush() 

205 

206 def _flush(self): 

207 value = self.buffer.getvalue() 

208 if value: 

209 self.flush() 

210 self.target.write(value) 

211 self.buffer.seek(0) 

212 self.buffer.truncate(0) 

213 self.needs_clear = False 

214 

215 

216class StreamWrapper(object): 

217 '''Wrap stdout and stderr globally''' 

218 

219 def __init__(self): 

220 self.stdout = self.original_stdout = sys.stdout 

221 self.stderr = self.original_stderr = sys.stderr 

222 self.original_excepthook = sys.excepthook 

223 self.wrapped_stdout = 0 

224 self.wrapped_stderr = 0 

225 self.wrapped_excepthook = 0 

226 self.capturing = 0 

227 self.listeners = set() 

228 

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

230 self.wrap_stdout() 

231 

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

233 self.wrap_stderr() 

234 

235 def start_capturing(self, bar=None): 

236 if bar: # pragma: no branch 

237 self.listeners.add(bar) 

238 

239 self.capturing += 1 

240 self.update_capturing() 

241 

242 def stop_capturing(self, bar=None): 

243 if bar: # pragma: no branch 

244 try: 

245 self.listeners.remove(bar) 

246 except KeyError: 

247 pass 

248 

249 self.capturing -= 1 

250 self.update_capturing() 

251 

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

253 if isinstance(self.stdout, WrappingIO): 

254 self.stdout.capturing = self.capturing > 0 

255 

256 if isinstance(self.stderr, WrappingIO): 

257 self.stderr.capturing = self.capturing > 0 

258 

259 if self.capturing <= 0: 

260 self.flush() 

261 

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

263 if stdout: 

264 self.wrap_stdout() 

265 

266 if stderr: 

267 self.wrap_stderr() 

268 

269 def wrap_stdout(self): 

270 self.wrap_excepthook() 

271 

272 if not self.wrapped_stdout: 

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

274 listeners=self.listeners) 

275 self.wrapped_stdout += 1 

276 

277 return sys.stdout 

278 

279 def wrap_stderr(self): 

280 self.wrap_excepthook() 

281 

282 if not self.wrapped_stderr: 

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

284 listeners=self.listeners) 

285 self.wrapped_stderr += 1 

286 

287 return sys.stderr 

288 

289 def unwrap_excepthook(self): 

290 if self.wrapped_excepthook: 

291 self.wrapped_excepthook -= 1 

292 sys.excepthook = self.original_excepthook 

293 

294 def wrap_excepthook(self): 

295 if not self.wrapped_excepthook: 

296 logger.debug('wrapping excepthook') 

297 self.wrapped_excepthook += 1 

298 sys.excepthook = self.excepthook 

299 

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

301 if stdout: 

302 self.unwrap_stdout() 

303 

304 if stderr: 

305 self.unwrap_stderr() 

306 

307 def unwrap_stdout(self): 

308 if self.wrapped_stdout > 1: 

309 self.wrapped_stdout -= 1 

310 else: 

311 sys.stdout = self.original_stdout 

312 self.wrapped_stdout = 0 

313 

314 def unwrap_stderr(self): 

315 if self.wrapped_stderr > 1: 

316 self.wrapped_stderr -= 1 

317 else: 

318 sys.stderr = self.original_stderr 

319 self.wrapped_stderr = 0 

320 

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

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

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

324 return stderr_needs_clear or stdout_needs_clear 

325 

326 def flush(self): 

327 if self.wrapped_stdout: # pragma: no branch 

328 try: 

329 self.stdout._flush() 

330 except (io.UnsupportedOperation, 

331 AttributeError): # pragma: no cover 

332 self.wrapped_stdout = False 

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

334 sys.stdout) 

335 

336 if self.wrapped_stderr: # pragma: no branch 

337 try: 

338 self.stderr._flush() 

339 except (io.UnsupportedOperation, 

340 AttributeError): # pragma: no cover 

341 self.wrapped_stderr = False 

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

343 sys.stderr) 

344 

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

346 self.original_excepthook(exc_type, exc_value, exc_traceback) 

347 self.flush() 

348 

349 

350class AttributeDict(dict): 

351 ''' 

352 A dict that can be accessed with .attribute 

353 

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

355 

356 # Reading 

357 >>> attrs['spam'] 

358 123 

359 >>> attrs.spam 

360 123 

361 

362 # Read after update using attribute 

363 >>> attrs.spam = 456 

364 >>> attrs['spam'] 

365 456 

366 >>> attrs.spam 

367 456 

368 

369 # Read after update using dict access 

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

371 >>> attrs['spam'] 

372 123 

373 >>> attrs.spam 

374 123 

375 

376 # Read after update using dict access 

377 >>> del attrs.spam 

378 >>> attrs['spam'] 

379 Traceback (most recent call last): 

380 ... 

381 KeyError: 'spam' 

382 >>> attrs.spam 

383 Traceback (most recent call last): 

384 ... 

385 AttributeError: No such attribute: spam 

386 >>> del attrs.spam 

387 Traceback (most recent call last): 

388 ... 

389 AttributeError: No such attribute: spam 

390 ''' 

391 def __getattr__(self, name): 

392 if name in self: 

393 return self[name] 

394 else: 

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

396 

397 def __setattr__(self, name, value): 

398 self[name] = value 

399 

400 def __delattr__(self, name): 

401 if name in self: 

402 del self[name] 

403 else: 

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

405 

406 

407logger = logging.getLogger(__name__) 

408streams = StreamWrapper()