diff --git a/nut/cli.py b/nut/cli.py index 25d2c7e..df7d4ec 100644 --- a/nut/cli.py +++ b/nut/cli.py @@ -1,10 +1,30 @@ import argparse +from typing import Any, Callable, Iterable, Optional + from colorama import Fore, Style, init init(autoreset=True) +_INTERNAL_ARG_FIELDS = frozenset({"command", "func", "_args", "help"}) +_COLOR_TOKEN_MAP = { + "%red%": Fore.RED, + "%green%": Fore.GREEN, + "%yellow%": Fore.YELLOW, + "%blue%": Fore.BLUE, + "%magenta%": Fore.MAGENTA, + "%cyan%": Fore.CYAN, + "%white%": Fore.WHITE, + "%reset%": Style.RESET_ALL, + "%bold%": Style.BRIGHT, +} + + class CLI: - def __init__(self, prog="nut", help_callback=None): + def __init__( + self, + prog: str = "nut", + help_callback: Optional[Callable[[], Iterable[str]]] = None, + ): self.parser = argparse.ArgumentParser(prog=prog, add_help=False) self.subparsers = self.parser.add_subparsers(dest="command") @@ -12,73 +32,74 @@ class CLI: self.parser.add_argument("-h", "--help", action="store_true") - def command(self, name): - parser = self.subparsers.add_parser(name, add_help=False) - cmd = _Command(parser) - return cmd + def command(self, name: str) -> "_Command": + command_parser = self.subparsers.add_parser(name, add_help=False) + return _Command(command_parser) - def run(self): - args = self.parser.parse_args() + def run(self) -> int: + parsed_args = self.parser.parse_args() - if getattr(args, "help", False) or not getattr(args, "command", None): - if self.help_callback: - for line in self.help_callback(): - print(self._color(line)) - else: - self.parser.print_help() + if self._should_show_help(parsed_args): + self._print_help() return 0 - args_list = getattr(args, "_args", []) + handler = getattr(parsed_args, "func", None) + if handler is None: + raise RuntimeError("No handler defined") + + args_list = list(getattr(parsed_args, "_args", [])) flags = { - k: v for k, v in vars(args).items() - if k not in ("command", "func", "_args", "help") + k: v for k, v in vars(parsed_args).items() + if k not in _INTERNAL_ARG_FIELDS } - return args.func(args_list, flags) + return handler(args_list, flags) + + def _should_show_help(self, parsed_args: argparse.Namespace) -> bool: + return bool( + getattr(parsed_args, "help", False) + or not getattr(parsed_args, "command", None) + ) + + def _print_help(self) -> None: + if self.help_callback is None: + self.parser.print_help() + return + + for line in self.help_callback(): + print(self._color(str(line))) def _color(self, text: str) -> str: - mapping = { - "%red%": Fore.RED, - "%green%": Fore.GREEN, - "%yellow%": Fore.YELLOW, - "%blue%": Fore.BLUE, - "%magenta%": Fore.MAGENTA, - "%cyan%": Fore.CYAN, - "%white%": Fore.WHITE, - "%reset%": Style.RESET_ALL, - "%bold%": Style.BRIGHT, - } - - for key, value in mapping.items(): + for key, value in _COLOR_TOKEN_MAP.items(): text = text.replace(key, value) return text + class _Command: - def __init__(self, parser): + def __init__(self, parser: argparse.ArgumentParser): self.parser = parser - self.parser.set_defaults(func=self._missing) + self.parser.set_defaults(func=self._missing_handler) self.parser.add_argument("_args", nargs="*") self.parser.add_argument("-h", "--help", action="store_true") - def arg(self, name, **kwargs): + def arg(self, name: str, **kwargs: Any) -> "_Command": self.parser.add_argument(name, **kwargs) return self - def flag(self, name, short=None): + def flag(self, name: str, short: Optional[str] = None) -> "_Command": opts = [f"--{name}"] if short: - opts.append(f"-{short}") + opts.append(f"-{short.lstrip('-')}") self.parser.add_argument(*opts, action="store_true") return self - def handle(self, func): + def handle(self, func: Callable[..., int]) -> "_Command": self.parser.set_defaults(func=func) return self - def _missing(self, *_): + def _missing_handler(self, *_: Any) -> int: raise RuntimeError("No handler defined") - \ No newline at end of file