Coverage for src/edwh/helpers.py: 22%

1"""

2This file contains re-usable helpers.

3"""

5import abc

6import datetime as dt

7import functools

8import io

9import itertools

10import os

11import re

12import sys

13import typing as t

14from pathlib import Path

16import click

17import diceware

18import invoke

19import yaml

20from ewok import Context

21from more_itertools import flatten as _flatten

23from .constants import DOCKER_COMPOSE, AnyDict

26def confirm(prompt: str, default: bool = False, allowed: set[str] | None = None, strict: bool = False) -> bool:

27 """

28 Prompt a user to confirm a (dangerous) action.

29 By default, entering nothing (only enter) will result in False, unless 'default' is set to True.

30 """

31 if os.environ.get("EDWH_NON_INTERACTIVE", "0") == "1": 31 ↛ 32line 31 didn't jump to line 32 because the condition on line 31 was never true

32 if strict:

33 raise RuntimeError(f"Prevented strict `confirm({prompt})` in --non-interactive mode")

34 else:

35 return default

37 allowed = allowed or {"y", "t", "1"}

38 if default: 38 ↛ 41line 38 didn't jump to line 41 because the condition on line 38 was always true

39 allowed.add(" ")

41 answer = input(prompt).lower().strip()

42 answer += " "

44 confirmed = answer.strip() in allowed or answer[0] in allowed

46 if strict and not confirmed: 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true

47 raise RuntimeError(f"Stopping now because '{answer.strip()}' did not match {allowed}.")

49 return confirmed

52def executes_correctly(c: Context, argument: str) -> bool:

53 """returns True if the execution was without error level"""

54 ran = c.run(argument, warn=True, hide=True)

55 return bool(ran and ran.ok)

58def execution_fails(c: Context, argument: str) -> bool:

59 """Returns true if the execution fails based on error level"""

60 return not executes_correctly(c, argument)

63def run_pty(ctx: Context, *command_parts: str, **options) -> invoke.Result | None:

64 try:

65 command = " ".join(command_parts)

66 return ctx.run(command, pty=True, **options)

67 except invoke.exceptions.Failure:

68 # error is already printed due to `pty`

69 return None

72def run_pty_ok(ctx: Context, *command_parts: str, **options) -> bool:

73 result = run_pty(ctx, *command_parts, **options)

74 return bool(result and result.ok)

77def generate_password(silent: bool = True, dice: int = 6) -> str:

78 """Generate a diceware password using --dice 6."""

79 options = diceware.handle_options(args=[])

80 options.num = dice

81 password: str = diceware.get_passphrase(options)

82 if not silent:

83 print("Password:", password)

84 return password

87def _add_dash(flag: str) -> str:

88 if flag.startswith("-"):

89 # don't change

90 return flag

91 if len(flag) == 1:

92 # one letter, -x

93 return f"-{flag}"

94 else:

95 # multiple letters --flag

96 return f"--{flag}"

99def arg_was_passed(flag: str | tuple[str, ...]) -> int | None:

100 """

101 Returns the index of the flag in sys.argv if passed, else None

102 """

103 flag = flag if isinstance(flag, tuple) else (flag,)

104 flag = tuple(_add_dash(f) for f in flag)

105 # flag and sys.argv should now both be in the same format: -x and --flag

106 return next((i for i, item in enumerate(sys.argv) if item in flag), None)

107

108

109def kwargs_to_options(data: AnyDict | None = None, **kw: t.Any) -> str:

110 """

111 Convert a dictionary of options to the cli variant

112 e.g. {'a': 1, 'key': 2} -> -a 1 --key 2

113 """

114 if data:

115 kw |= data

116

117 options = []

118 for key, value in kw.items():

119 if value in (None, "", False):

120 # skip falsey, but keep 0

121 continue

122

123 pref = ("-" if len(key) == 1 else "--") + key

124

125 if isinstance(value, bool):

126 options.append(f"{pref}")

127

128 elif isinstance(value, list):

129 options.extend(f"{pref} {subvalue}" for subvalue in value)

130 else:

131 options.append(f"{pref} {value}")

132

133 return " " + " ".join(options)

134

135

136class Logger(abc.ABC):

137 def log(self, *a: t.Any) -> None:

138 raise NotImplementedError("This is an abstract method")

139

140

141class VerboseLogger(Logger):

142 def __init__(self) -> None:

143 self._then = self._now()

144 self._previous = self._now()

145

146 @staticmethod

147 def _now() -> dt.datetime:

148 return dt.datetime.now(dt.timezone.utc)

149

150 def log(self, *a: t.Any) -> None:

151 now = self._now()

152 delta_start = now - self._then

153 delta_prev = now - self._previous

154 print(f"[{delta_start}, +{delta_prev}]", *a, file=sys.stderr)

155 self._previous = now

156

157

158# usage:

159# logger = VerboseLogger()

160# log = logger.log

161# ...

162# log("some event")

163

164

165class NoopLogger(Logger):

166 def log(self, *_: t.Any) -> None:

167 return None

168

169

170def noop(*_: t.Any, **__: t.Any) -> None:

171 return None

172

173

174@t.overload

175def dump_set_as_list[T](data: set[T]) -> list[T]:

176 """

177 Sets are converted to lists.

178 """

179

180

181@t.overload

182def dump_set_as_list[T](data: T) -> T:

183 """

184 Other datatypes remain untouched.

185 """

186

187

188def dump_set_as_list[T](data: set[T] | T) -> list[T] | T:

189 if isinstance(data, set):

190 return list(data)

191 else:

192 return data

193

194

195KEY_ENTER = "\r"

196KEY_ARROWUP = "\033[A"

197KEY_ARROWDOWN = "\033[B"

198

199

200def print_box(label: str, selected: bool, current: bool, number: int, fmt: str = "[%s]", filler: str = "x") -> None:

201 box = fmt % (filler if selected else " ")

202 indicator = ">" if current else " "

203 click.echo(f"{indicator}{number}. {box} {label}")

204

205

206def interactive_selected_checkbox_values[H: t.Hashable](

207 options: list[str] | dict[H, str],

208 prompt: str = "Select options (use arrow keys, spacebar, or digit keys, press 'Enter' to finish):",

209 selected: t.Collection[H] = (),

210 allow_empty: bool = False,

211) -> list[str | H] | None:

212 """

213 This function provides an interactive checkbox selection in the console.

214

215 The user can navigate through the options using the arrow keys,

216 select/deselect options using the spacebar or digit keys, and finish the selection by pressing 'Enter'.

217

218 Args:

219 options: A list or dict (value: label) of options to be displayed as checkboxes.

220 prompt (str, optional): A string that is displayed as a prompt for the user.

221 allow_empty (bool, optional): If True, adds an extra option "(none)" to deselect all other options.

222 selected: a set (/other iterable) of pre-selected options (set is preferred).

223

224 `H: Hashable` means the values have to be the same type as the keys of options

225 and they should be hashable (via `hash()`).

226 Example:

227 options = {1: "something", "two": "else"}

228 selected = [2, "three"] # valid type (int and str are keys of options)

229 selected = [1.5, "two"] # invalid type (none of the keys of options are a float)

230

231 Returns:

232 list[str]: A list of selected option values.

233

234 Examples:

235 interactive_selected_checkbox_values(["first", "second", "third"])

236

237 interactive_selected_checkbox_values({100: "first", 211: "second", 355: "third"})

238

239 interactive_selected_checkbox_values(["first", "second", "third"], selected=["third"])

240

241 interactive_selected_checkbox_values({1: "first", 2: "second", 3: "third"}, selected=[3])

242 """

243 checked_indices: dict[int, str | H] = {} # instead of set to keep ordering

244 current_index = 0

245

246 if isinstance(options, list):

247 labels = options

248 option_values = t.cast(t.Sequence[str | H], options)

249 else:

250 labels = list(options.values())

251 option_values = t.cast(t.Sequence[str | H], list(options))

252

253 for item in selected:

254 if item not in option_values:

255 # invalid

256 continue

257

258 idx = option_values.index(item)

259 checked_indices[idx] = option_values[idx]

260

261 if allow_empty:

262 labels.append("(none)")

263

264 print_checkbox = functools.partial(print_box, fmt="[%s]", filler="x")

265

266 while True:

267 click.clear()

268 click.echo(prompt)

269

270 for i, option in enumerate(labels, start=1):

271 print_checkbox(option, i - 1 in checked_indices, i - 1 == current_index, i)

272

273 key = click.getchar()

274

275 if key == KEY_ENTER:

276 break

277 elif key == KEY_ARROWUP: # Up arrow

278 current_index = (current_index - 1) % len(labels)

279 elif key == KEY_ARROWDOWN: # Down arrow

280 current_index = (current_index + 1) % len(labels)

281 elif key.isdigit() and 1 <= int(key) <= len(labels):

282 current_index = int(key) - 1

283 elif key == " ":

284 if allow_empty and current_index == len(labels) - 1:

285 checked_indices.clear()

286 checked_indices[len(labels) - 1] = "(none)"

287 else:

288 if len(checked_indices) == 1 and set(checked_indices.values()) == {"(none)"}:

289 checked_indices.clear()

290 if current_index in checked_indices:

291 del checked_indices[current_index]

292 else:

293 checked_indices[current_index] = option_values[current_index]

294

295 if allow_empty and len(checked_indices) == 1 and set(checked_indices.values()) == {"(none)"}:

296 # None instead of empty list since otherwise it would just ask again

297 return None

298

299 return list(checked_indices.values())

300

301

302def interactive_selected_radio_value[H: t.Hashable](

303 options: list[str] | dict[H, str],

304 prompt: str = "Select an option (use arrow keys, spacebar, or digit keys, press 'Enter' to finish):",

305 selected: H | None = None,

306 allow_empty: bool = False,

307) -> str | H | None:

308 """

309 This function provides an interactive radio box selection in the console.

310

311 The user can navigate through the options using the arrow keys,

312 select an option using the spacebar or digit keys, and finish the selection by pressing 'Enter'.

313

314 Args:

315 options: A list or dict (value: label) of options to be displayed as radio boxes.

316 prompt (str, optional): A string that is displayed as a prompt for the user.

317 allow_empty (bool, optional): If True, adds an extra option "(none)" to allow deselecting all options.

318 selected: a pre-selected option.

319 `H: Hashable` means the values have to be the same type as the keys of options

320 and they should be hashable (via `hash()`).

321 Example:

322 options = {1: "something", "two": "else"}

323 selected = 2 # valid type (int is a key of options)

324 selected = 1.5 # invalid type (none of the keys of options are a float)

325

326 Returns:

327 str: The selected option value, or an empty string if (none) is selected.

328

329 Examples:

330 interactive_selected_radio_value(["first", "second", "third"])

331

332 interactive_selected_radio_value({100: "first", 211: "second", 355: "third"})

333

334 interactive_selected_radio_value(["first", "second", "third"], selected="third")

335

336 interactive_selected_radio_value({1: "first", 2: "second", 3: "third"}, selected=3)

337 """

338 selected_index: int | None = None

339 current_index = 0

340

341 if isinstance(options, list):

342 labels = options

343 option_values = t.cast(t.Sequence[str | H], options)

344 else:

345 labels = list(options.values())

346 option_values = t.cast(t.Sequence[str | H], list(options))

347

348 if selected in option_values:

349 selected_index = current_index = option_values.index(selected)

350

351 if allow_empty:

352 labels.append("(none)")

353

354 print_radio_box = functools.partial(print_box, fmt="(%s)", filler="o")

355

356 while True:

357 click.clear()

358 click.echo(prompt)

359

360 for i, option in enumerate(labels, start=1):

361 print_radio_box(option, i - 1 == selected_index, i - 1 == current_index, i)

362

363 key = click.getchar()

364

365 if key == KEY_ENTER:

366 if selected_index is None:

367 # no you may not leave.

368 continue

369 else:

370 # done!

371 break

372

373 elif key == KEY_ARROWUP: # Up arrow

374 current_index = (current_index - 1) % len(labels)

375 elif key == KEY_ARROWDOWN: # Down arrow

376 current_index = (current_index + 1) % len(labels)

377 elif key.isdigit() and 1 <= int(key) <= len(labels):

378 selected_index = int(key) - 1

379 elif key == " ":

380 selected_index = current_index

381

382 if allow_empty and selected_index == len(labels) - 1:

383 return None

384

385 return option_values[selected_index]

386

387

388def yaml_loads(text: str) -> AnyDict:

389 dct = yaml.load(

390 text,

391 Loader=yaml.SafeLoader,

392 )

393 return t.cast(AnyDict, dct)

394

395

396def dc_config(ctx: Context) -> AnyDict:

397 if ran := ctx.run(f"{DOCKER_COMPOSE} config", warn=True, echo=False, hide=True):

398 return (

399 yaml_loads(

400 ran.stdout.strip(),

401 )

402 or {}

403 )

404 else:

405 return {}

406

407

408def print_aligned(plugin_commands: list[str]) -> None:

409 """

410 Prints a list of plugin commands in an aligned format.

411

412 This function takes a list of plugin commands, each of which is a string containing two parts separated by a tab.

413 It splits each command into two parts, calculates the maximum length of the first part across all commands,

414 and then prints each command with the first part left-justified to the maximum length. This ensures that the

415 second parts of all commands are aligned in the output.

416

417 Args:

418 plugin_commands (list[str]): A list of plugin commands. Each command is a string containing two parts

419 separated by a tab.

420

421 Example:

422 print_aligned(["command1\tdescription1", "command_with_long_name\tdescription2"])

423 # Output:

424 # command1 description1

425 # command_with_long_name description2

426 """

427 splitted = [_.split("\t") for _ in plugin_commands]

428 max_l = max([len(_[0]) for _ in splitted])

429

430 for before, after in splitted:

431 print("\t", before.ljust(max_l, " "), "\t\t", after)

432

433

434def flatten[T](something: t.Iterable[t.Iterable[T]]) -> list[T]:

435 """

436 Like itertools.flatten but eager

437 """

438 return list(_flatten(something))

439

440

441def shorten(text: str, max_chars: int) -> str:

442 """

443 textwrap looks at words and stuff, not relevant for commands!

444 """

445 if len(text) <= max_chars:

446 return text

447 else:

448 return f"{text[:max_chars]}..."

449

450

451def _fabric_resolve_home(path: str, user: str) -> str:

452 if not path.startswith("~"):

453 return path

454

455 return path.replace("~", f"/home/{user}", 1)

456

457

458def _write_bytes_remote(c: Context, path: str, contents: bytes, parents: bool = False) -> None:

459 f = io.BytesIO(contents)

460

461 if parents:

462 # ensure path to file exists

463 parent_path = Path(path).parent

464 c.run(f"mkdir -p {parent_path}")

465

466 c.put(f, path)

467

468

469def _write_bytes_local(_: Context, path: str, contents: bytes, parents: bool = False) -> None:

470 p = Path(path)

471 if parents:

472 p.parent.mkdir(parents=True, exist_ok=True)

473

474 p.write_bytes(contents)

475

476

477class WriteBytesFn(t.Protocol):

478 def __call__(self, c: Context, path: str, contents: bytes, parents: bool = False) -> None: ... 478 ↛ exitline 478 didn't return from function '__call__' because

479

480

481def fabric_write(c: Context, path: str, contents: str | bytes, parents: bool = False) -> None:

482 """

483 Write some contents to a remote file.

484 ~ will be resolved to the remote user's home

485 """

486 path = _fabric_resolve_home(path, c.user) if c.user else path

487

488 fn = t.cast(WriteBytesFn, _write_bytes_local if isinstance(c, invoke.Context) else _write_bytes_remote)

489

490 return fn(c, path, contents if isinstance(contents, bytes) else contents.encode(), parents=parents)

491

492

493def _read_bytes_remote(c: Context, path: str) -> bytes:

494 buf = io.BytesIO()

495 c.get(path, buf)

496

497 buf.seek(0)

498 return buf.read()

499

500

501def _read_bytes_local(_: Context, path: str) -> bytes:

502 return Path(path).read_bytes()

503

504

505type ReadBytesFn = t.Callable[[Context, str], bytes]

506

507

508def fabric_read_bytes(c: Context, path: str, throw: bool = True) -> bytes:

509 """

510 Write some bytes from a remote file.

511 ~ will be resolved to the remote user's home

512 """

513 path = _fabric_resolve_home(path, c.user) if c.user else path

514

515 fn: ReadBytesFn = _read_bytes_local if isinstance(c, invoke.Context) else _read_bytes_remote

516

517 try:

518 return fn(c, path)

519 except FileNotFoundError:

520 if throw:

521 raise

522 else:

523 return b""

524

525

526def fabric_read(c: Context, path: str, throw: bool = True) -> str:

527 """

528 Write some text from a remote file.

529 ~ will be resolved to the remote user's home

530 """

531 b = fabric_read_bytes(c, path, throw=throw)

532 return b.decode()

533

534

535def _add_alias(sometask: t.Any, alias: str):

536 if alias not in sometask.aliases:

537 sometask.aliases = (*sometask.aliases, alias)

538

539

540def add_alias(sometask: t.Any, aliases: str | t.Iterable[str]):

541 """

542 Add an extra alias to an existing task (usually in ~/.config/edwh/tasks.py).

543

544 Example:

545 >>> edwh.add_alias(edwh.tasks.migrate, "migarte")

546 >>> edwh.add_alias(edwh.tasks.migrate, ["migarte"])

547 >>> edwh.add_alias(edwh.tasks.migrate, ("migarte",))

548 """

549 if isinstance(aliases, str):

550 aliases = [aliases]

551

552 for alias in aliases:

553 _add_alias(sometask, alias)

554

555

556type ColorFn = t.Callable[[str], str]

557type FilterFn = t.Callable[[str], bool] | None

558

559

560class Handler(abc.ABC):

561 @abc.abstractmethod

562 def process(self, chunk: str):

563 pass

564

565

566class NoopHandler(Handler):

567 def process(self, chunk: str): # noqa: ARG002

568 return

569

570

571class LineBufferHandler(Handler):

572 def __init__(self, prefix: str, output_stream: t.TextIO, filter_fn: FilterFn = None):

573 self.buffer = ""

574 self.prefix = prefix

575 self.output_stream = output_stream

576 self.filter_fn = filter_fn

577

578 def process(self, chunk: str):

579 if not chunk:

580 return

581

582 filter_fn = self.filter_fn

583 self.buffer += chunk

584

585 if "\n" in self.buffer:

586 lines = self.buffer.splitlines(True)

587

588 if not self.buffer.endswith("\n"):

589 # last line wasn't finished yet, save it to buffer instead of printing

590 self.buffer = lines.pop()

591 else:

592 # all lines complete, clean buffer

593 self.buffer = ""

594

595 for line in lines:

596 if filter_fn and not filter_fn(line):

597 continue

598

599 print(f"{self.prefix}{line}", end="", flush=True, file=self.output_stream)

600

601

602POSSIBLE_FLAGS = {

603 # https://docs.python.org/3/library/re.html

604 "a": re.ASCII,

605 "d": re.DEBUG,

606 "i": re.IGNORECASE,

607 "l": re.LOCALE,

608 "m": None, # re.MULTILINE but the logger works line-by-line so this isn't really possible

609 "s": re.DOTALL,

610 "u": re.UNICODE,

611 "x": re.VERBOSE,

612 # custom: 'v' to invert

613}

614

615

616def parse_regex(raw: str) -> FilterFn:

617 """

618 Turn `/pattern/flags` into a Regex object.

619

620 Uses the grep style flags (i for case insensitive, v for invert)

621 """

622

623 # zero slashes: just a pattern, no flags.

624 # one slash: search term with / in it

625 # two slashes (+ starts with /): regex with flags

626 # more slashes: flags AND / in filter itself

627

628 if raw.startswith("/") and raw.count("/") > 1:

629 # flag-mode

630 _, *rest, flags_str = raw.split("/")

631 flags = set(flags_str.lower())

632 pattern = "/".join(rest)

633 else:

634 # normal search mode, no flags

635 flags = set()

636 pattern = raw

637

638 flags_bin = 0 # re.NOFLAG doesn't exist in 3.10 yet

639

640 for flag in flags:

641 flags_bin |= POSSIBLE_FLAGS.get(flag) or 0 # re.NOFLAG

642

643 re_compiled = re.compile(pattern, flags_bin)

644

645 if "v" in flags:

646 # v for inverse like `grep -v`

647 return lambda text: not re_compiled.search(text)

648 else:

649 return lambda text: bool(re_compiled.search(text))

650

651

652def ansi_color_code(code: str, format_opts: t.Collection[str] = ()) -> str:

653 res = "\033["

654 for c in format_opts:

655 res += f"{c};"

656 return f"{res}{code}m"

657

658

659def make_color_func(code: str) -> ColorFn:

660 return lambda s: f"{ansi_color_code(code)}{s}{ansi_color_code('0')}"

661

662

663def build_rainbow() -> tuple[ColorFn, ...]:

664 names = (

665 "grey",

666 "red",

667 "green",

668 "yellow",

669 "blue",

670 "magenta",

671 "cyan",

672 "white",

673 )

674

675 colors = {}

676 for i, name in enumerate(names):

677 colors[name] = make_color_func(str(30 + i))

678 colors[f"intense_{name}"] = make_color_func(f"{30 + i};1")

679

680 return (

681 colors["cyan"],

682 colors["yellow"],

683 colors["green"],

684 colors["magenta"],

685 colors["blue"],

686 colors["intense_cyan"],

687 colors["intense_yellow"],

688 colors["intense_green"],

689 colors["intense_magenta"],

690 colors["intense_blue"],

691 )

692

693

694def rainbow() -> t.Generator[ColorFn, None, None]:

695 """

696 rainbow = []colorFunc{

697 colors["cyan"],

698 colors["yellow"],

699 colors["green"],

700 colors["magenta"],

701 colors["blue"],

702 colors["intense_cyan"],

703 colors["intense_yellow"],

704 colors["intense_green"],

705 colors["intense_magenta"],

706 colors["intense_blue"],

707 }

708

709 Yield colors from the docker compose rainbow map in a cyclic way.

710 """

711 yield from itertools.cycle(build_rainbow())