Merge pull request #1881 from fabito/feature/exclude-args-from-ui

Ability to hide extra args from web ui

examples/add_command_line_argument.py

@@ -5,6 +5,8 @@
 @events.init_command_line_parser.add_listener
 def _(parser):
     parser.add_argument("--my-argument", type=str, env_var="LOCUST_MY_ARGUMENT", default="", help="It's working")
+    # Set `include_in_web_ui` to False if you want to hide from the web UI
+    parser.add_argument("--my-ui-invisible-argument", include_in_web_ui=False, default="I am invisible")
 
 
 @events.init.add_listener
@@ -18,4 +20,5 @@ class WebsiteUser(HttpUser):
 
     @task
     def my_task(self):
-        print(self.environment.parsed_options.my_argument)
+        print(f"my_argument={self.environment.parsed_options.my_argument}")
+        print(f"my_ui_invisible_argument={self.environment.parsed_options.my_ui_invisible_argument}")

locust/argument_parser.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import textwrap
+from typing import Dict
 
 import configargparse
 
@@ -13,6 +14,32 @@
 DEFAULT_CONFIG_FILES = ["~/.locust.conf", "locust.conf"]
 
 
+class LocustArgumentParser(configargparse.ArgumentParser):
+    """Drop-in replacement for `configargparse.ArgumentParser` that adds support for
+    optionally exclude arguments from the UI.
+    """
+
+    def add_argument(self, *args, **kwargs) -> configargparse.Action:
+        """
+        This method supports the same args as ArgumentParser.add_argument(..)
+        as well as the additional args below.
+
+        Arguments:
+            include_in_web_ui: If True (default), the argument will show in the UI.
+
+        Returns:
+            argparse.Action: the new argparse action
+        """
+        include_in_web_ui = kwargs.pop("include_in_web_ui", True)
+        action = super().add_argument(*args, **kwargs)
+        action.include_in_web_ui = include_in_web_ui
+        return action
+
+    @property
+    def args_included_in_web_ui(self) -> Dict[str, configargparse.Action]:
+        return {a.dest: a for a in self._actions if hasattr(a, "include_in_web_ui") and a.include_in_web_ui}
+
+
 def _is_package(path):
     """
     Is the given path a Python package?
@@ -55,7 +82,7 @@ def find_locustfile(locustfile):
 
 
 def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG_FILES):
-    parser = configargparse.ArgumentParser(
+    parser = LocustArgumentParser(
         default_config_files=default_config_files,
         add_env_var_help=False,
         add_config_file_help=False,
@@ -470,3 +497,14 @@ def default_args_dict():
     default_parser = get_empty_argument_parser()
     setup_parser_arguments(default_parser)
     return vars(default_parser.parse([]))
+
+
+def ui_extra_args_dict(args=None) -> Dict[str, str]:
+    """Get all the UI visible arguments"""
+    locust_args = default_args_dict()
+
+    parser = get_parser()
+    all_args = vars(parser.parse_args(args))
+
+    extra_args = {k: v for k, v in all_args.items() if k not in locust_args and k in parser.args_included_in_web_ui}
+    return extra_args

locust/test/test_parser.py

@@ -5,7 +5,7 @@
 from io import StringIO
 
 import locust
-from locust.argument_parser import parse_options, get_parser, parse_locustfile_option
+from locust.argument_parser import parse_options, get_parser, parse_locustfile_option, ui_extra_args_dict
 from .mock_locustfile import mock_locustfile
 from .testcases import LocustTestCase
 
@@ -174,3 +174,27 @@ def test_csv_full_history_requires_csv(self):
                         "--csv-full-history",
                     ]
                 )
+
+    def test_custom_argument_included_in_web_ui(self):
+        @locust.events.init_command_line_parser.add_listener
+        def _(parser, **kw):
+            parser.add_argument("--a1", help="a1 help")
+            parser.add_argument("--a2", help="a2 help", include_in_web_ui=False)
+
+        args = [
+            "-u",
+            "666",
+            "--a1",
+            "v1",
+            "--a2",
+            "v2",
+        ]
+        options = parse_options(args=args)
+        self.assertEqual(666, options.num_users)
+        self.assertEqual("v1", options.a1)
+        self.assertEqual("v2", options.a2)
+
+        extra_args = ui_extra_args_dict(args)
+        self.assertIn("a1", extra_args)
+        self.assertNotIn("a2", extra_args)
+        self.assertEqual("v1", extra_args["a1"])

locust/web.py

@@ -416,15 +416,7 @@ def update_template_args(self):
             "total": get_task_ratio_dict(self.environment.user_classes, total=True),
         }
 
-        extra_options = (
-            {
-                k: v
-                for k, v in vars(self.environment.parsed_options).items()
-                if k not in argument_parser.default_args_dict()
-            }
-            if self.environment.parsed_options
-            else {}
-        )
+        extra_options = argument_parser.ui_extra_args_dict()
 
         self.template_args = {
             "locustfile": self.environment.locustfile,