Refactor CLI class: enhance type hints, improve help handling, and streamline argument parsing

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")
-