[PATCH v4 3/3] dts: rework arguments framework

[Luca Vizzarro <luca.vizzarro@arm.com>]

The existing argument handling in the code relies on basic argparse
functionality and a custom argparse action to integrate environment
variables. This commit improves the current handling by augmenting
argparse.

This rework implements the following improvements:
- There are duplicate expressions scattered throughout the code. To
  improve readability and maintainability, these are refactored
  into list/dict comprehensions or factory functions.
- Instead of relying solely on argument flags, error messages now
  accurately reference environment variables when applicable, enhancing
  user experience. For instance:

    error: environment variable DTS_DPDK_TARBALL: Invalid file

- Knowing the number of environment variables and arguments set
  allow for a useful help page display when none are provided.
- A display of which environment variables have been detected and their
  corresponding values in the help page, aiding user awareness.

Signed-off-by: Luca Vizzarro <luca.vizzarro@arm.com>
---
 doc/guides/tools/dts.rst  |  53 +++++---
 dts/framework/settings.py | 280 +++++++++++++++++++++++++++-----------
 2 files changed, 235 insertions(+), 98 deletions(-)

diff --git a/doc/guides/tools/dts.rst b/doc/guides/tools/dts.rst
index 47b218b2c6..6993443389 100644
--- a/doc/guides/tools/dts.rst
+++ b/doc/guides/tools/dts.rst
@@ -215,30 +215,41 @@ DTS is run with ``main.py`` located in the ``dts`` directory after entering Poet
 .. code-block:: console
 
    (dts-py3.10) $ ./main.py --help
-   usage: main.py [-h] [--config-file CONFIG_FILE] [--output-dir OUTPUT_DIR] [-t TIMEOUT] [-v] [-s] [--tarball TARBALL] [--compile-timeout COMPILE_TIMEOUT] [--test-suite TEST_SUITE [TEST_CASES ...]] [--re-run RE_RUN]
+   usage: main.py [-h] [--config-file FILE_PATH] [--output-dir DIR_PATH] [-t SECONDS] [-v] [-s] [--tarball FILE_PATH]
+                  [--compile-timeout SECONDS] [--test-suite TEST_SUITE [TEST_CASES ...]] [--re-run N_TIMES]
 
-   Run DPDK test suites. All options may be specified with the environment variables provided in brackets. Command line arguments have higher priority.
+   Run DPDK test suites. All options may be specified with the environment variables provided in brackets. Command
+   line arguments have higher priority.
 
    options:
-   -h, --help            show this help message and exit
-   --config-file CONFIG_FILE
-                         [DTS_CFG_FILE] The configuration file that describes the test cases, SUTs and targets. (default: ./conf.yaml)
-   --output-dir OUTPUT_DIR, --output OUTPUT_DIR
-                         [DTS_OUTPUT_DIR] Output directory where DTS logs and results are saved. (default: output)
-   -t TIMEOUT, --timeout TIMEOUT
-                         [DTS_TIMEOUT] The default timeout for all DTS operations except for compiling DPDK. (default: 15)
-   -v, --verbose         [DTS_VERBOSE] Specify to enable verbose output, logging all messages to the console. (default: False)
-   -s, --skip-setup      [DTS_SKIP_SETUP] Specify to skip all setup steps on SUT and TG nodes. (default: False)
-   --tarball TARBALL, --snapshot TARBALL, --git-ref TARBALL
-                         [DTS_DPDK_TARBALL] Path to DPDK source code tarball or a git commit ID, tag ID or tree ID to test. To test local changes, first commit them, then use the commit ID with this option. (default: dpdk.tar.xz)
-   --compile-timeout COMPILE_TIMEOUT
-                         [DTS_COMPILE_TIMEOUT] The timeout for compiling DPDK. (default: 1200)
-   --test-suite TEST_SUITE [TEST_CASES ...]
-                         [DTS_TEST_SUITES] A list containing a test suite with test cases. The first parameter is the test suite name, and the rest are test case names, which are optional. May be specified multiple times. To specify multiple test suites in the environment
-                         variable, join the lists with a comma. Examples: --test-suite suite case case --test-suite suite case ... | DTS_TEST_SUITES='suite case case, suite case, ...' | --test-suite suite --test-suite suite case ... | DTS_TEST_SUITES='suite, suite case, ...'
-                         (default: [])
-   --re-run RE_RUN, --re_run RE_RUN
-                         [DTS_RERUN] Re-run each test case the specified number of times if a test failure occurs. (default: 0)
+     -h, --help            show this help message and exit
+     --config-file FILE_PATH
+                           [DTS_CFG_FILE] The configuration file that describes the test cases, SUTs and targets.
+                           (default: conf.yaml)
+     --output-dir DIR_PATH, --output DIR_PATH
+                           [DTS_OUTPUT_DIR] Output directory where DTS logs and results are saved. (default: output)
+     -t SECONDS, --timeout SECONDS
+                           [DTS_TIMEOUT] The default timeout for all DTS operations except for compiling DPDK.
+                           (default: 15)
+     -v, --verbose         [DTS_VERBOSE] Specify to enable verbose output, logging all messages to the console.
+                           (default: False)
+     -s, --skip-setup      [DTS_SKIP_SETUP] Specify to skip all setup steps on SUT and TG nodes. (default: False)
+     --tarball FILE_PATH, --snapshot FILE_PATH, --git-ref FILE_PATH
+                           [DTS_DPDK_TARBALL] Path to DPDK source code tarball or a git commit ID,tag ID or tree ID to
+                           test. To test local changes, first commit them, then use the commit ID with this option.
+                           (default: dpdk.tar.xz)
+     --compile-timeout SECONDS
+                           [DTS_COMPILE_TIMEOUT] The timeout for compiling DPDK. (default: 1200)
+     --test-suite TEST_SUITE [TEST_CASES ...]
+                           [DTS_TEST_SUITES] A list containing a test suite with test cases. The first parameter is
+                           the test suite name, and the rest are test case names, which are optional. May be specified
+                           multiple times. To specify multiple test suites in the environment variable, join the lists
+                           with a comma. Examples: --test-suite SUITE1 CASE1 CASE2 --test-suite SUITE2 CASE1 ... |
+                           DTS_TEST_SUITES='SUITE1 CASE1 CASE2, SUITE2 CASE1, ...' | --test-suite SUITE1 --test-suite SUITE2
+                           CASE1 ... | DTS_TEST_SUITES='SUITE1, SUITE2 CASE1, ...' (default: [])
+     --re-run N_TIMES, --re_run N_TIMES
+                           [DTS_RERUN] Re-run each test case the specified number of times if a test failure occurs.
+                           (default: 0)
 
 
 The brackets contain the names of environment variables that set the same thing.
diff --git a/dts/framework/settings.py b/dts/framework/settings.py
index 688e8679a7..b19f274f9d 100644
--- a/dts/framework/settings.py
+++ b/dts/framework/settings.py
@@ -2,6 +2,7 @@
 # Copyright(c) 2010-2021 Intel Corporation
 # Copyright(c) 2022-2023 PANTHEON.tech s.r.o.
 # Copyright(c) 2022 University of New Hampshire
+# Copyright(c) 2024 Arm Limited
 
 """Environment variables and command line arguments parsing.
 
@@ -72,9 +73,15 @@
 
 import argparse
 import os
+import sys
+from argparse import (
+    Action,
+    ArgumentDefaultsHelpFormatter,
+    _get_action_name,
+)
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
+from typing import Callable, ParamSpec
 
 from .config import TestSuiteConfig
 from .utils import DPDKGitTarball
@@ -109,104 +116,224 @@ class Settings:
 
 SETTINGS: Settings = Settings()
 
+P = ParamSpec("P")
 
-def _get_parser() -> argparse.ArgumentParser:
-    """Create the argument parser for DTS.
+#: Attribute name representing the env variable name to augment :class:`~argparse.Action` with.
+_ENV_VAR_NAME_ATTR = "env_var_name"
+#: Attribute name representing the value origin to augment :class:`~argparse.Action` with.
+_IS_FROM_ENV_ATTR = "is_from_env"
 
-    Command line options take precedence over environment variables, which in turn take precedence
-    over default values.
+#: The prefix to be added to all of the environment variables.
+ENV_PREFIX = "DTS_"
+
+
+def is_action_in_args(action: Action) -> bool:
+    """Check if the action is invoked in the command line arguments."""
+    for option in action.option_strings:
+        if option in sys.argv:
+            return True
+    return False
+
+
+def make_env_var_name(action: Action, env_var_name: str | None) -> str:
+    """Make and assign an environment variable name to the given action."""
+    env_var_name = f"{ENV_PREFIX}{env_var_name or action.dest.upper()}"
+    setattr(action, _ENV_VAR_NAME_ATTR, env_var_name)
+    return env_var_name
+
+
+def get_env_var_name(action: Action) -> str | None:
+    """Get the environment variable name of the given action."""
+    return getattr(action, _ENV_VAR_NAME_ATTR, None)
+
+
+def set_is_from_env(action: Action) -> None:
+    """Make the environment the given action's value origin."""
+    setattr(action, _IS_FROM_ENV_ATTR, True)
 
-    Returns:
-        argparse.ArgumentParser: The configured argument parser with defined options.
-    """
 
-    def env_arg(env_var: str, default: Any) -> Any:
-        """A helper function augmenting the argparse with environment variables.
+def is_from_env(action: Action) -> bool:
+    """Check if the given action's value originated from the environment."""
+    return getattr(action, _IS_FROM_ENV_ATTR, False)
 
-        If the supplied environment variable is defined, then the default value
-        of the argument is modified. This satisfies the priority order of
-        command line argument > environment variable > default value.
 
-        Args:
-            env_var: Environment variable name.
-            default: Default value.
+def augment_add_argument_with_env(
+    add_argument_fn: Callable[P, Action],
+):
+    """Augment any :class:`~argparse._ActionsContainer.add_argument` with environment variables."""
 
-        Returns:
-            Environment variable or default value.
+    def _add_argument(
+        *args: P.args,
+        env_var_name: str | None = None,
+        **kwargs: P.kwargs,
+    ) -> Action:
+        """Add an argument with an environment variable to the parser."""
+        action = add_argument_fn(*args, **kwargs)
+        env_var_name = make_env_var_name(action, env_var_name)
+
+        if not is_action_in_args(action):
+            env_var_value = os.environ.get(env_var_name)
+            if env_var_value:
+                set_is_from_env(action)
+                sys.argv[1:0] = [action.format_usage(), env_var_value]
+
+        return action
+
+    return _add_argument
+
+
+class ArgumentParser(argparse.ArgumentParser):
+    """ArgumentParser with a custom error message.
+
+    This custom version of ArgumentParser changes the error message to
+    accurately reflect its origin if an environment variable is used
+    as an argument.
+
+    Instead of printing usage on every error, it prints instructions
+    on how to do it.
+    """
+
+    def find_action(
+        self, action_name: str, filter_fn: Callable[[Action], bool] | None = None
+    ) -> Action | None:
+        """Find and return an action by its name.
+
+        Arguments:
+            action_name: the name of the action to find.
+            filter_fn: a filter function to use in the search.
         """
-        return os.environ.get(env_var) or default
+        it = (action for action in self._actions if action_name == _get_action_name(action))
+        action = next(it, None)
+
+        if action is not None and filter_fn is not None:
+            return action if filter_fn(action) else None
+
+        return action
 
-    parser = argparse.ArgumentParser(
+    def error(self, message):
+        """Augments :meth:`~argparse.ArgumentParser.error` with environment variable awareness."""
+        for action in self._actions:
+            if is_from_env(action):
+                action_name = _get_action_name(action)
+                env_var_name = get_env_var_name(action)
+                env_var_value = os.environ.get(env_var_name)
+
+                message = message.replace(
+                    f"argument {action_name}",
+                    f"environment variable {env_var_name} (value: {env_var_value})",
+                )
+
+        print(f"{self.prog}: error: {message}\n", file=sys.stderr)
+        self.exit(2, "For help and usage, " "run the command with the --help flag.\n")
+
+
+class EnvVarHelpFormatter(ArgumentDefaultsHelpFormatter):
+    """Custom formatter to add environment variables in the help page."""
+
+    def _get_help_string(self, action):
+        """Overrides :meth:`ArgumentDefaultsHelpFormatter._get_help_string`."""
+        help = super()._get_help_string(action)
+
+        env_var_name = get_env_var_name(action)
+        if env_var_name is not None:
+            help = f"[{env_var_name}] {help}"
+
+            env_var_value = os.environ.get(env_var_name)
+            if env_var_value is not None:
+                help += f" (env value: {env_var_value})"
+
+        return help
+
+
+def _get_parser() -> ArgumentParser:
+    """Create the argument parser for DTS.
+
+    Command line options take precedence over environment variables, which in turn take precedence
+    over default values.
+
+    Returns:
+        ArgumentParser: The configured argument parser with defined options.
+    """
+    parser = ArgumentParser(
         description="Run DPDK test suites. All options may be specified with the environment "
         "variables provided in brackets. Command line arguments have higher priority.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        formatter_class=EnvVarHelpFormatter,
+        allow_abbrev=False,
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env = augment_add_argument_with_env(parser.add_argument)
+
+    add_argument_to_parser_with_env(
         "--config-file",
-        default=env_arg("DTS_CFG_FILE", SETTINGS.config_file_path),
+        default=SETTINGS.config_file_path,
         type=Path,
-        help="[DTS_CFG_FILE] The configuration file that describes the test cases, "
-        "SUTs and targets.",
+        help="The configuration file that describes the test cases, SUTs and targets.",
+        metavar="FILE_PATH",
+        env_var_name="CFG_FILE",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "--output-dir",
         "--output",
-        default=env_arg("DTS_OUTPUT_DIR", SETTINGS.output_dir),
-        help="[DTS_OUTPUT_DIR] Output directory where DTS logs and results are saved.",
+        default=SETTINGS.output_dir,
+        help="Output directory where DTS logs and results are saved.",
+        metavar="DIR_PATH",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "-t",
         "--timeout",
-        default=env_arg("DTS_TIMEOUT", SETTINGS.timeout),
+        default=SETTINGS.timeout,
         type=float,
-        help="[DTS_TIMEOUT] The default timeout for all DTS operations except for compiling DPDK.",
+        help="The default timeout for all DTS operations except for compiling DPDK.",
+        metavar="SECONDS",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "-v",
         "--verbose",
         action="store_true",
-        default=env_arg("DTS_VERBOSE", SETTINGS.verbose),
-        help="[DTS_VERBOSE] Specify to enable verbose output, logging all messages "
-        "to the console.",
+        default=SETTINGS.verbose,
+        help="Specify to enable verbose output, logging all messages to the console.",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "-s",
         "--skip-setup",
         action="store_true",
-        default=env_arg("DTS_SKIP_SETUP", SETTINGS.skip_setup),
-        help="[DTS_SKIP_SETUP] Specify to skip all setup steps on SUT and TG nodes.",
+        default=SETTINGS.skip_setup,
+        help="Specify to skip all setup steps on SUT and TG nodes.",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "--tarball",
         "--snapshot",
         "--git-ref",
-        default=env_arg("DTS_DPDK_TARBALL", SETTINGS.dpdk_tarball_path),
+        default=SETTINGS.dpdk_tarball_path,
         type=Path,
-        help="[DTS_DPDK_TARBALL] Path to DPDK source code tarball or a git commit ID, "
+        help="Path to DPDK source code tarball or a git commit ID, "
         "tag ID or tree ID to test. To test local changes, first commit them, "
         "then use the commit ID with this option.",
+        metavar="FILE_PATH",
+        dest="dpdk_tarball_path",
+        env_var_name="DPDK_TARBALL",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "--compile-timeout",
-        default=env_arg("DTS_COMPILE_TIMEOUT", SETTINGS.compile_timeout),
+        default=SETTINGS.compile_timeout,
         type=float,
-        help="[DTS_COMPILE_TIMEOUT] The timeout for compiling DPDK.",
+        help="The timeout for compiling DPDK.",
+        metavar="SECONDS",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "--test-suite",
         action="append",
         nargs="+",
         metavar=("TEST_SUITE", "TEST_CASES"),
-        default=env_arg("DTS_TEST_SUITES", SETTINGS.test_suites),
-        help="[DTS_TEST_SUITES] A list containing a test suite with test cases. "
+        default=SETTINGS.test_suites,
+        help="A list containing a test suite with test cases. "
         "The first parameter is the test suite name, and the rest are test case names, "
         "which are optional. May be specified multiple times. To specify multiple test suites in "
         "the environment variable, join the lists with a comma. "
@@ -215,21 +342,23 @@ def env_arg(env_var: str, default: Any) -> Any:
         "DTS_TEST_SUITES='suite case case, suite case, ...' | "
         "--test-suite suite --test-suite suite case ... | "
         "DTS_TEST_SUITES='suite, suite case, ...'",
+        dest="test_suites",
     )
 
-    parser.add_argument(
+    add_argument_to_parser_with_env(
         "--re-run",
         "--re_run",
-        default=env_arg("DTS_RERUN", SETTINGS.re_run),
+        default=SETTINGS.re_run,
         type=int,
-        help="[DTS_RERUN] Re-run each test case the specified number of times "
-        "if a test failure occurs.",
+        help="Re-run each test case the specified number of times if a test failure occurs.",
+        env_var_name="RERUN",
+        metavar="N_TIMES",
     )
 
     return parser
 
 
-def _process_test_suites(args: str | list[list[str]]) -> list[TestSuiteConfig]:
+def _process_test_suites(parser: ArgumentParser, args: list[list[str]]) -> list[TestSuiteConfig]:
     """Process the given argument to a list of :class:`TestSuiteConfig` to execute.
 
     Args:
@@ -240,17 +369,12 @@ def _process_test_suites(args: str | list[list[str]]) -> list[TestSuiteConfig]:
     Returns:
         A list of test suite configurations to execute.
     """
-    if isinstance(args, str):
-        # Environment variable in the form of "suite case case, suite case, suite, ..."
-        args = [suite_with_cases.split() for suite_with_cases in args.split(",")]
+    test_suites = parser.find_action("test_suites", is_from_env)
+    if test_suites is not None:
+        # Environment variable in the form of "SUITE1 CASE1 CASE2, SUITE2 CASE1, SUITE3, ..."
+        args = [suite_with_cases.split() for suite_with_cases in args[0][0].split(",")]
 
-    test_suites_to_run = []
-    for suite_with_cases in args:
-        test_suites_to_run.append(
-            TestSuiteConfig(test_suite=suite_with_cases[0], test_cases=suite_with_cases[1:])
-        )
-
-    return test_suites_to_run
+    return [TestSuiteConfig(test_suite, test_cases) for [test_suite, *test_cases] in args]
 
 
 def get_settings() -> Settings:
@@ -261,19 +385,21 @@ def get_settings() -> Settings:
     Returns:
         The new settings object.
     """
-    parsed_args = _get_parser().parse_args()
-    return Settings(
-        config_file_path=parsed_args.config_file,
-        output_dir=parsed_args.output_dir,
-        timeout=parsed_args.timeout,
-        verbose=parsed_args.verbose,
-        skip_setup=parsed_args.skip_setup,
-        dpdk_tarball_path=Path(
-            Path(DPDKGitTarball(parsed_args.tarball, parsed_args.output_dir))
-            if not os.path.exists(parsed_args.tarball)
-            else Path(parsed_args.tarball)
-        ),
-        compile_timeout=parsed_args.compile_timeout,
-        test_suites=_process_test_suites(parsed_args.test_suite),
-        re_run=parsed_args.re_run,
+    parser = _get_parser()
+
+    if len(sys.argv) == 1:
+        parser.print_help()
+        sys.exit(1)
+
+    args = parser.parse_args()
+
+    args.dpdk_tarball_path = Path(
+        Path(DPDKGitTarball(args.dpdk_tarball_path, args.output_dir))
+        if not os.path.exists(args.dpdk_tarball_path)
+        else Path(args.dpdk_tarball_path)
     )
+
+    args.test_suites = _process_test_suites(parser, args.test_suites)
+
+    kwargs = {k: v for k, v in vars(args).items() if hasattr(SETTINGS, k)}
+    return Settings(**kwargs)
-- 
2.34.1