Coverage for gws-app/gws/lib/cli/__init__.py: 66%

1"""Utilities for command line scripts. 

2 

3This package is used by the ``gws`` command line commands and by the 

4development and test scripts. It has no dependencies on the rest of GWS and 

5can be imported outside of the GWS application. 

6 

7It provides: 

8 

9- colored console output: ``cprint``, ``info``, ``warning``, ``error`` and 

10 ``fatal``, which exits the script. Messages are prefixed with 

11 ``SCRIPT_NAME`` when set, colors are used only on a terminal. 

12- shell commands: ``run`` echoes and runs a command and exits on failure, 

13 ``exec`` runs a command and returns its output. 

14- simple file utilities: ``find_dirs``, ``find_files``, ``ensure_dir``, 

15 ``read_file``, ``write_file``. 

16- a script entry point: ``parse_args`` parses ``-opt value`` and 

17 ``--opt value`` style arguments, ``main`` runs a main function with them 

18 and prints the usage on ``-h``. 

19- ``text_table``, which formats rows as a plain text table. 

20- ``ProgressIndicator``, a context manager that logs the progress of a 

21 long-running task in percent steps. 

22 

23Example:: 

24 

25 import gws.lib.cli as cli 

26 

27 USAGE = ''' 

28 Usage: myscript.py <dir> [-pattern <regex>] 

29 ''' 

30 

31 def main(args): 

32 paths = list(cli.find_files(args[1], args.get('pattern'))) 

33 cli.info(cli.text_table([{'path': p} for p in paths], header='auto')) 

34 with cli.ProgressIndicator('processing', len(paths)) as pi: 

35 for p in paths: 

36 pi.update() 

37 return 0 

38 

39 if __name__ == '__main__': 

40 cli.main('myscript', main, USAGE) 

41""" 

42 

43import re 

44import os 

45import sys 

46import subprocess 

47import time 

48import math 

49import traceback 

50 

51SCRIPT_NAME = '' 

52"""Name of the running script, used as a prefix for messages.""" 

53 

54_COLOR = { 

55 'black': '\x1b[30m', 

56 'red': '\x1b[31m', 

57 'green': '\x1b[32m', 

58 'yellow': '\x1b[33m', 

59 'blue': '\x1b[34m', 

60 'magenta': '\x1b[35m', 

61 'cyan': '\x1b[36m', 

62 'white': '\x1b[37m', 

63 'reset': '\x1b[0m', 

64} 

65 

66 

67def cprint(clr, msg): 

68 """Print a message to stdout, in color if stdout is a terminal. 

69 

70 Args: 

71 clr: Color name, e.g. ``red`` or ``cyan``, or an empty value for no color. 

72 msg: Message. 

73 """ 

74 if SCRIPT_NAME: 

75 msg = '[' + SCRIPT_NAME + '] ' + msg 

76 if clr and sys.stdout.isatty(): 

77 msg = _COLOR[clr] + msg + _COLOR['reset'] 

78 sys.stdout.write(msg + '\n') 

79 sys.stdout.flush() 

80 

81 

82def error(msg): 

83 """Print an error message in red. 

84 

85 Args: 

86 msg: Message. 

87 """ 

88 cprint('red', msg) 

89 

90 

91def fatal(msg): 

92 """Print an error message in red and exit with code 1. 

93 

94 Args: 

95 msg: Message. 

96 """ 

97 cprint('red', msg) 

98 sys.exit(1) 

99 

100 

101def warning(msg): 

102 """Print a warning message in yellow. 

103 

104 Args: 

105 msg: Message. 

106 """ 

107 cprint('yellow', msg) 

108 

109 

110def info(msg): 

111 """Print an info message in cyan. 

112 

113 Args: 

114 msg: Message. 

115 """ 

116 cprint('cyan', msg) 

117 

118 

119## 

120 

121def run(cmd): 

122 """Print and run a shell command, exit the script if it fails. 

123 

124 The command output is not captured. 

125 

126 Args: 

127 cmd: Command as a string or a list of strings, which are joined with spaces. 

128 """ 

129 if isinstance(cmd, list): 

130 cmd = ' '.join(cmd) 

131 cmd = re.sub(r'\s+', ' ', cmd.strip()) 

132 info(f'> {cmd}') 

133 res = subprocess.run(cmd, shell=True, capture_output=False) 

134 if res.returncode: 

135 fatal(f'COMMAND FAILED, code {res.returncode}') 

136 

137 

138def exec(cmd): 

139 """Run a shell command and return its output. 

140 

141 Args: 

142 cmd: Command string. 

143 

144 Returns: 

145 The stripped stdout of the command, or an error message if the command could not be run. 

146 """ 

147 try: 

148 return ( 

149 subprocess 

150 .run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) 

151 .stdout.decode('utf8').strip() 

152 ) 

153 except Exception as exc: 

154 return f'> {cmd} FAILED: {exc}' 

155 

156 

157def find_dirs(dirname): 

158 """Find the subdirectories of a directory, not recursively. 

159 

160 Hidden directories are skipped. 

161 

162 Args: 

163 dirname: Directory path. 

164 

165 Yields: 

166 Subdirectory paths. Nothing if ``dirname`` is not a directory. 

167 """ 

168 if not os.path.isdir(dirname): 

169 return 

170 

171 de: os.DirEntry 

172 for de in os.scandir(dirname): 

173 if de.name.startswith('.'): 

174 continue 

175 if de.is_dir(): 

176 yield de.path 

177 

178 

179def find_files(dirname, pattern=None, deep=True): 

180 """Find files in a directory. 

181 

182 Hidden files and directories are skipped. 

183 

184 Args: 

185 dirname: Directory path. 

186 pattern: Regular expression to search for in file paths. 

187 deep: Search subdirectories too. 

188 

189 Yields: 

190 File paths. Nothing if ``dirname`` is not a directory. 

191 """ 

192 if not os.path.isdir(dirname): 

193 return 

194 

195 de: os.DirEntry 

196 for de in os.scandir(dirname): 

197 if de.name.startswith('.'): 

198 continue 

199 if de.is_dir() and deep: 

200 yield from find_files(de.path, pattern) 

201 continue 

202 if de.is_file() and (pattern is None or re.search(pattern, de.path)): 

203 yield de.path 

204 

205 

206def ensure_dir(path): 

207 """Create a directory, including parent directories. 

208 

209 Args: 

210 path: Directory path. 

211 

212 Returns: 

213 The path. 

214 """ 

215 os.makedirs(path, exist_ok=True) 

216 return path 

217 

218 

219def read_file(path): 

220 """Read a text file. 

221 

222 Args: 

223 path: File path. 

224 

225 Returns: 

226 The file content, stripped. 

227 """ 

228 with open(path, 'rt', encoding='utf8') as fp: 

229 return fp.read().strip() 

230 

231 

232def write_file(path, text): 

233 """Write a text file. 

234 

235 Args: 

236 path: File path. 

237 text: Content. 

238 """ 

239 with open(path, 'wt', encoding='utf8') as fp: 

240 fp.write(text) 

241 

242 

243def parse_args(argv): 

244 """Parse command line arguments. 

245 

246 ``-opt`` and ``--opt`` set the option ``opt`` to ``True``, a following non-option argument 

247 sets it to that value instead. Other arguments are stored under integer keys, in order. 

248 A ``-`` argument stores all remaining arguments as a list under ``_rest``. 

249 

250 Args: 

251 argv: Arguments, usually ``sys.argv``. 

252 

253 Returns: 

254 A dict of options and positional arguments. 

255 """ 

256 args = {} 

257 opt = None 

258 n = 0 

259 

260 for a in argv: 

261 if a == '-': 

262 args['_rest'] = [] 

263 elif '_rest' in args: 

264 args['_rest'].append(a) 

265 elif a.startswith('--'): 

266 opt = a[2:] 

267 args[opt] = True 

268 elif a.startswith('-'): 

269 opt = a[1:] 

270 args[opt] = True 

271 elif opt: 

272 args[opt] = a 

273 opt = None 

274 else: 

275 args[n] = a 

276 n += 1 

277 

278 return args 

279 

280 

281def main(name, main_fn, usage): 

282 """Run the main function of a script. 

283 

284 Parses ``sys.argv`` and calls ``main_fn`` with the parsed arguments. With ``-h`` or ``--help``, 

285 prints the usage text and exits. The return value of ``main_fn`` is used as the exit code. 

286 Exceptions are printed as internal errors and exit with code 1, keyboard interrupts exit with code 130. 

287 

288 Args: 

289 name: Script name, used as a prefix for messages. 

290 main_fn: Main function, called with the dict from ``parse_args``. 

291 usage: Usage text. 

292 """ 

293 global SCRIPT_NAME 

294 

295 SCRIPT_NAME = name 

296 

297 args = parse_args(sys.argv) 

298 if not args or 'h' in args or 'help' in args: 

299 print('\n' + usage.strip() + '\n') 

300 sys.exit(0) 

301 

302 try: 

303 sys.exit(main_fn(args)) 

304 except KeyboardInterrupt: 

305 sys.exit(130) 

306 except Exception as exc: 

307 error('INTERNAL ERROR') 

308 error(traceback.format_exc()) 

309 sys.exit(1) 

310 

311 

312def text_table(data, header=None, delim=' | '): 

313 """Format rows as a plain text table. 

314 

315 Numbers are right-aligned, other values left-aligned. 

316 

317 Args: 

318 data: Rows, either dicts or sequences. 

319 header: Column keys, or ``auto`` to use the keys (or indexes) of the first row. 

320 If given, a header line is printed. If ``None``, the columns of the first row are used without a header. 

321 delim: Column delimiter. 

322 

323 Returns: 

324 The table text, or an empty string if there are no rows. 

325 """ 

326 

327 data = list(data) 

328 

329 if not data: 

330 return '' 

331 

332 is_dict = isinstance(data[0], dict) 

333 

334 print_header = header is not None 

335 if header is None or header == 'auto': 

336 header = data[0].keys() if is_dict else list(range(len(data[0]))) 

337 

338 widths = [len(h) if print_header else 1 for h in header] 

339 

340 def get(d, h): 

341 if is_dict: 

342 return d.get(h, '') 

343 try: 

344 return d[h] 

345 except IndexError: 

346 return '' 

347 

348 for d in data: 

349 widths = [ 

350 max(a, b) 

351 for a, b in zip( 

352 widths, 

353 [len(str(get(d, h))) for h in header] 

354 ) 

355 ] 

356 

357 def field(n, v): 

358 if isinstance(v, (int, float)): 

359 return str(v).rjust(widths[n]) 

360 return str(v).ljust(widths[n]) 

361 

362 rows = [] 

363 

364 if print_header: 

365 hdr = delim.join(field(n, h) for n, h in enumerate(header)) 

366 rows.append(hdr) 

367 rows.append('-' * len(hdr)) 

368 

369 for d in data: 

370 rows.append(delim.join(field(n, get(d, h)) for n, h in enumerate(header))) 

371 

372 return '\n'.join(rows) 

373 

374 

375class ProgressIndicator: 

376 """Context manager that logs the progress of a task. 

377 

378 Logs ``START`` on enter, the progress in percent steps on ``update``, 

379 and ``END`` with the elapsed time on a normal exit. 

380 """ 

381 

382 def __init__(self, title, total=0, resolution=10): 

383 """Create a progress indicator. 

384 

385 Args: 

386 title: Title, used as a prefix for messages. 

387 total: Total number of items. If 0, no progress is logged. 

388 resolution: Step in percent between progress messages. 

389 """ 

390 self.resolution = resolution 

391 self.title = title 

392 self.total = total 

393 self.progress = 0 

394 self.lastd = 0 

395 self.starttime = 0 

396 

397 def __enter__(self): 

398 self.log(f'START ({self.total})' if self.total else 'START') 

399 self.starttime = time.time() 

400 return self 

401 

402 def __exit__(self, exc_type, exc_val, exc_tb): 

403 if not exc_type: 

404 ts = time.time() - self.starttime 

405 self.log(f'END ({ts:.2f} sec)') 

406 

407 def update(self, add=1): 

408 """Add processed items and log the progress if it reached the next step. 

409 

410 Args: 

411 add: Number of processed items. 

412 """ 

413 if not self.total: 

414 return 

415 self.progress += add 

416 p = math.floor(self.progress * 100.0 / self.total) 

417 if p > 100: 

418 p = 100 

419 d = round(p / self.resolution) * self.resolution 

420 if d > self.lastd: 

421 self.log(f'{d}%') 

422 self.lastd = d 

423 

424 def log(self, s): 

425 """Log a message with the title. 

426 

427 Args: 

428 s: Message. 

429 """ 

430 info(f'{self.title}: {s}')