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

nut/cli.py

@@ -1,44 +1,12 @@
 import argparse
+from typing import Any, Callable, Iterable, Optional
+
 from colorama import Fore, Style, init
 
 init(autoreset=True)
 
-class CLI:
+_INTERNAL_ARG_FIELDS = frozenset({"command", "func", "_args", "help"})
-    def __init__(self, prog="nut", help_callback=None):
+_COLOR_TOKEN_MAP = {
-        self.parser = argparse.ArgumentParser(prog=prog, add_help=False)
-        self.subparsers = self.parser.add_subparsers(dest="command")
-
-        self.help_callback = help_callback
-
-        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 run(self):
-        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()
-            return 0
-
-        args_list = getattr(args, "_args", [])
-
-        flags = {
-            k: v for k, v in vars(args).items()
-            if k not in ("command", "func", "_args", "help")
-        }
-
-        return args.func(args_list, flags)
-
-    def _color(self, text: str) -> str:
-        mapping = {
     "%red%": Fore.RED,
     "%green%": Fore.GREEN,
     "%yellow%": Fore.YELLOW,
@@ -48,37 +16,90 @@ class CLI:
     "%white%": Fore.WHITE,
     "%reset%": Style.RESET_ALL,
     "%bold%": Style.BRIGHT,
+}
+
+
+class CLI:
+    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")
+
+        self.help_callback = help_callback
+
+        self.parser.add_argument("-h", "--help", action="store_true")
+
+    def command(self, name: str) -> "_Command":
+        command_parser = self.subparsers.add_parser(name, add_help=False)
+        return _Command(command_parser)
+
+    def run(self) -> int:
+        parsed_args = self.parser.parse_args()
+
+        if self._should_show_help(parsed_args):
+            self._print_help()
+            return 0
+
+        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(parsed_args).items()
+            if k not in _INTERNAL_ARG_FIELDS
         }
 
-        for key, value in mapping.items():
+        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:
+        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")
-