diff --git a/examples/add_command_line_argument.py b/examples/add_command_line_argument.py index 1b23ffc70a..69136f36dc 100644 --- a/examples/add_command_line_argument.py +++ b/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}") diff --git a/locust/argument_parser.py b/locust/argument_parser.py index c4887908fe..8b539689df 100644 --- a/locust/argument_parser.py +++ b/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 diff --git a/locust/test/test_parser.py b/locust/test/test_parser.py index 1999eb3cc6..e674c7d61f 100644 --- a/locust/test/test_parser.py +++ b/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"]) diff --git a/locust/web.py b/locust/web.py index 68cb6594b7..6de01e8057 100644 --- a/locust/web.py +++ b/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,