From c230f7a1c2000f9b15fd54f99c82d487e06fe734 Mon Sep 17 00:00:00 2001
From: Renan Rodrigo <renanrodrigo@canonical.com>
Date: Fri, 12 Jul 2024 14:03:54 -0300
Subject: [PATCH] cli: add method to get help for any command directly from
 parser

And the parser is now ProArgumentParser

Signed-off-by: Renan Rodrigo <renanrodrigo@canonical.com>
---
 uaclient/cli/__init__.py | 14 ++++++++++++--
 uaclient/cli/config.py   | 34 ++++++----------------------------
 uaclient/cli/help.py     |  2 +-
 uaclient/cli/system.py   |  2 +-
 4 files changed, 20 insertions(+), 32 deletions(-)

diff --git a/uaclient/cli/__init__.py b/uaclient/cli/__init__.py
index 8eed3ccc81..82aa474324 100644
--- a/uaclient/cli/__init__.py
+++ b/uaclient/cli/__init__.py
@@ -55,7 +55,7 @@
 ]
 
 
-class UAArgumentParser(argparse.ArgumentParser):
+class ProArgumentParser(argparse.ArgumentParser):
     def error(self, message):
         self.print_usage(sys.stderr)
         # In some cases (e.g. `pro --wrong-flag`) argparse errors out asking
@@ -72,9 +72,19 @@ def error(self, message):
             message = messages.CLI_TRY_HELP
         self.exit(2, message + "\n")
 
+    def print_help_for_command(self, command: str):
+        args_list = command.split()
+        args_list.append("--help")
+        try:
+            self.parse_args(args_list)
+        # We want help for any specific command,
+        # but without exiting right after
+        except SystemExit:
+            pass
+
 
 def get_parser(cfg: UAConfig):
-    parser = UAArgumentParser(
+    parser = ProArgumentParser(
         prog=NAME,
         formatter_class=argparse.RawDescriptionHelpFormatter,
         usage=USAGE_TMPL.format(name=NAME, command="<command>"),
diff --git a/uaclient/cli/config.py b/uaclient/cli/config.py
index 0c5a7af737..75cc54b1b9 100644
--- a/uaclient/cli/config.py
+++ b/uaclient/cli/config.py
@@ -25,9 +25,7 @@ def action_config(args, *, cfg, **kwargs):
     # Avoiding a circular import
     from uaclient.cli import get_parser
 
-    # Pretty sure there will be some method to do this soon
-    # so we don't repeat it all over the place
-    get_parser(cfg).parse_args(["config", "--help"])
+    get_parser(cfg).print_help_for_command("config")
     return 0
 
 
@@ -74,29 +72,17 @@ def action_config_set(args, *, cfg, **kwargs):
     try:
         set_key, set_value = args.key_value_pair.split("=")
     except ValueError:
-        try:
-            parser.parse_args(["config", "set", "--help"])
-        # Excepting SystemExit so we can raise our own exception
-        except SystemExit:
-            pass
+        parser.print_help_for_command("config set")
         raise exceptions.GenericInvalidFormat(
             expected="<key>=<value>", actual=args.key_value_pair
         )
     if set_key not in config.UA_CONFIGURABLE_KEYS:
-        try:
-            parser.parse_args(["config", "set", "--help"])
-        # Excepting SystemExit so we can raise our own exception
-        except SystemExit:
-            pass
+        parser.print_help_for_command("config set")
         raise exceptions.InvalidArgChoice(
             arg="<key>", choices=", ".join(config.UA_CONFIGURABLE_KEYS)
         )
     if not set_value.strip():
-        try:
-            parser.parse_args(["config", "set", "--help"])
-        # Excepting SystemExit so we can raise our own exception
-        except SystemExit:
-            pass
+        parser.print_help_for_command("config set")
         raise exceptions.EmptyConfigValue(arg=set_key)
     if set_key in ("http_proxy", "https_proxy"):
         protocol_type = set_key.split("_")[0]
@@ -182,11 +168,7 @@ def action_config_set(args, *, cfg, **kwargs):
             if set_value < 0:
                 raise ValueError("Invalid interval for {}".format(set_key))
         except ValueError:
-            try:
-                parser.parse_args(["config", "set", "--help"])
-            # Excepting SystemExit so we can raise our own exception
-            except SystemExit:
-                pass
+            parser.print_help_for_command("config set")
             # More readable in the CLI, without breaking the line in the logs
             print("")
             raise exceptions.InvalidPosIntConfigValue(
@@ -212,11 +194,7 @@ def action_config_unset(args, *, cfg, **kwargs):
 
     if args.key not in config.UA_CONFIGURABLE_KEYS:
         parser = get_parser(cfg=cfg)
-        try:
-            parser.parse_args(["config", "unset", "--help"])
-        # Excepting SystemExit so we can raise our own exception
-        except SystemExit:
-            pass
+        parser.print_help_for_command("config unset")
         raise exceptions.InvalidArgChoice(
             arg="<key>", choices=", ".join(config.UA_CONFIGURABLE_KEYS)
         )
diff --git a/uaclient/cli/help.py b/uaclient/cli/help.py
index 56031764a9..7ae1b9f18b 100644
--- a/uaclient/cli/help.py
+++ b/uaclient/cli/help.py
@@ -8,7 +8,7 @@ def action_help(args, *, cfg, **kwargs):
     service = args.service
 
     if not service:
-        # inserting it here to avoid circular imports
+        # Avoiding a circular import
         from uaclient.cli import get_parser
 
         get_parser(cfg=cfg).print_help()
diff --git a/uaclient/cli/system.py b/uaclient/cli/system.py
index 35ce1cdc00..626b595838 100644
--- a/uaclient/cli/system.py
+++ b/uaclient/cli/system.py
@@ -17,7 +17,7 @@ def action_system(args, *, cfg, **kwargs):
     # Avoiding a circular import
     from uaclient.cli import get_parser
 
-    get_parser(cfg).parse_args(["system", "--help"])
+    get_parser(cfg).print_help_for_command("system")
 
 
 reboot_required_subcommand = ProCommand(