Coverage for gws-app/gws/lib/cli/__init_

1"""Utilities for command line scripts.

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.

7It provides:

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.

23Example::

25 import gws.lib.cli as cli

27 USAGE = '''

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

29 '''

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

39 if __name__ == '__main__':

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

41"""

43import re

44import os

45import shutil

46import sys

47import subprocess

48import time

49import math

50import traceback

52SCRIPT_NAME = ''

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

55_COLOR = {

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

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

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

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

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

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

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

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

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

65}

68def cprint(clr, msg):

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

71 Args:

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

73 msg: Message.

74 """

75 if SCRIPT_NAME:

76 msg = '[' + SCRIPT_NAME + '] ' + msg

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

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

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

80 sys.stdout.flush()

83def error(msg):

84 """Print an error message in red.

86 Args:

87 msg: Message.

88 """

89 cprint('red', msg)

92def fatal(msg):

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

95 Args:

96 msg: Message.

97 """

98 cprint('red', msg)

99 sys.exit(1)

100

101

102def warning(msg):

103 """Print a warning message in yellow.

104

105 Args:

106 msg: Message.

107 """

108 cprint('yellow', msg)

109

110

111def info(msg):

112 """Print an info message in cyan.

113

114 Args:

115 msg: Message.

116 """

117 cprint('cyan', msg)

118

119

120##

121

122def run(cmd):

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

124

125 The command output is not captured.

126

127 Args:

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

129 """

130 if isinstance(cmd, list):

131 cmd = ' '.join(cmd)

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

133 info(f'> {cmd}')

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

135 if res.returncode:

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

137

138

139def exec(cmd):

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

141

142 Args:

143 cmd: Command string.

144

145 Returns:

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

147 """

148 try:

149 return (

150 subprocess

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

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

153 )

154 except Exception as exc:

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

156

157

158def find_dirs(dirname):

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

160

161 Hidden directories are skipped.

162

163 Args:

164 dirname: Directory path.

165

166 Yields:

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

168 """

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

170 return

171

172 de: os.DirEntry

173 for de in os.scandir(dirname):

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

175 continue

176 if de.is_dir():

177 yield de.path

178

179

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

181 """Find files in a directory.

182

183 Hidden files and directories are skipped.

184

185 Args:

186 dirname: Directory path.

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

188 deep: Search subdirectories too.

189

190 Yields:

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

192 """

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

194 return

195

196 de: os.DirEntry

197 for de in os.scandir(dirname):

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

199 continue

200 if de.is_dir() and deep:

201 yield from find_files(de.path, pattern)

202 continue

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

204 yield de.path

205

206

207def ensure_dir(path, clear=False):

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

209

210 Args:

211 path: Directory path.

212 clear: Remove the directory tree after creating it. The directory does not exist afterwards.

213

214 Returns:

215 The path.

216 """

217 os.makedirs(path, exist_ok=True)

218 if clear:

219 shutil.rmtree(path)

220 return path

221

222

223def read_file(path):

224 """Read a text file.

225

226 Args:

227 path: File path.

228

229 Returns:

230 The file content, stripped.

231 """

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

233 return fp.read().strip()

234

235

236def write_file(path, text):

237 """Write a text file.

238

239 Args:

240 path: File path.

241 text: Content.

242 """

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

244 fp.write(text)

245

246

247def parse_args(argv):

248 """Parse command line arguments.

249

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

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

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

253

254 Args:

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

256

257 Returns:

258 A dict of options and positional arguments.

259 """

260 args = {}

261 opt = None

262 n = 0

263

264 for a in argv:

265 if a == '-':

266 args['_rest'] = []

267 elif '_rest' in args:

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

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

270 opt = a[2:]

271 args[opt] = True

272 elif a.startswith('-'):

273 opt = a[1:]

274 args[opt] = True

275 elif opt:

276 args[opt] = a

277 opt = None

278 else:

279 args[n] = a

280 n += 1

281

282 return args

283

284

285def main(name, main_fn, usage):

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

287

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

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

290 Exceptions are printed as internal errors, keyboard interrupts are ignored.

291

292 Args:

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

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

295 usage: Usage text.

296 """

297 global SCRIPT_NAME

298

299 SCRIPT_NAME = name

300

301 args = parse_args(sys.argv)

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

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

304 sys.exit(0)

305

306 try:

307 sys.exit(main_fn(args))

308 except KeyboardInterrupt:

309 pass

310 except Exception as exc:

311 error('INTERNAL ERROR')

312 error(traceback.format_exc())

313

314

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

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

317

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

319

320 Args:

321 data: Rows, either dicts or sequences.

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

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

324 delim: Column delimiter.

325

326 Returns:

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

328 """

329

330 data = list(data)

331

332 if not data:

333 return ''

334

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

336

337 print_header = header is not None

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

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

340

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

342

343 def get(d, h):

344 if is_dict:

345 return d.get(h, '')

346 try:

347 return d[h]

348 except IndexError:

349 return ''

350

351 for d in data:

352 widths = [

353 max(a, b)

354 for a, b in zip(

355 widths,

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

357 )

358 ]

359

360 def field(n, v):

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

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

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

364

365 rows = []

366

367 if print_header:

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

369 rows.append(hdr)

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

371

372 for d in data:

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

374

375 return '\n'.join(rows)

376

377

378class ProgressIndicator:

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

380

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

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

383 """

384

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

386 """Create a progress indicator.

387

388 Args:

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

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

391 resolution: Step in percent between progress messages.

392 """

393 self.resolution = resolution

394 self.title = title

395 self.total = total

396 self.progress = 0

397 self.lastd = 0

398 self.starttime = 0

399

400 def __enter__(self):

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

402 self.starttime = time.time()

403 return self

404

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

406 if not exc_type:

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

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

409

410 def update(self, add=1):

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

412

413 Args:

414 add: Number of processed items.

415 """

416 if not self.total:

417 return

418 self.progress += add

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

420 if p > 100:

421 p = 100

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

423 if d > self.lastd:

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

425 self.lastd = d

426

427 def log(self, s):

428 """Log a message with the title.

429

430 Args:

431 s: Message.

432 """

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

Coverage for gws-app/gws/lib/cli/init.py: 65%

163 statements