Merge pull request #2612 from alzex3/master

Implement pyproject.toml support for Locust configuration

docs/configuration.rst

@@ -24,7 +24,8 @@ Locust is configured mainly through command line arguments.
 Environment Variables
 =====================
 
-Options can also be set through through environment variables. They are typically the same as the command line argument but capitalized and prefixed with ``LOCUST_``:
+Options can also be set through environment variables. They are typically the same as the command line argument
+but capitalized and prefixed with ``LOCUST_``:
 
 On Linux/macOS:
 
@@ -44,17 +45,23 @@ On Windows:
 Configuration File
 ==================
 
-Options can also be set in a configuration file in the `config file <https://github.com/bw2/ConfigArgParse#config-file-syntax>`_
-format. 
+Options can also be set in a configuration file in the
+`config or TOML file format <https://github.com/bw2/ConfigArgParse#config-file-syntax>`_.
 
-Locust will look for ``~/.locust.conf`` and ``./locust.conf`` by default, and you can specify an 
-additional file using the ``--config`` flag.
+Locust will look for ``~/.locust.conf``, ``./locust.conf`` and ``./pyproject.toml`` by default.
+You can specify an additional file using the ``--config`` flag.
 
-Example:
+.. code-block:: console
 
-.. code-block::
+    $ locust --config custom_config.conf
+
+Here's a quick example of the configuration files supported by Locust:
+
+locust.conf
+--------------
+
+.. code-block:: ini
 
-    # master.conf in current directory
     locustfile = locust_files/my_locust_file.py
     headless = true
     master = true
@@ -63,19 +70,34 @@ Example:
     users = 100
     spawn-rate = 10
     run-time = 10m
-    
+    tags = [Critical, Normal]
 
-.. code-block:: console
+pyproject.toml
+--------------
 
-    $ locust --config master.conf
+When using a TOML file, configuration options should be defined within the ``[tool.locust]`` section.
+
+.. code-block:: toml
+
+    [tool.locust]
+    locustfile = "locust_files/my_locust_file.py"
+    headless = true
+    master = true
+    expect-workers = 5
+    host = "https://target-system"
+    users = 100
+    spawn-rate = 10
+    run-time = "10m"
+    tags = ["Critical", "Normal"]
 
 .. note::
 
-    Configuration values are read (overridden) in the following order:
+    Configuration values are read (and overridden) in the following order:
 
     .. code-block:: console
         
-        ~/locust.conf -> ./locust.conf -> (file specified using --conf) -> env vars -> cmd args
+       ~/.locust.conf -> ./locust.conf -> ./pyproject.toml -> (file specified using --conf) -> env vars -> cmd args
+
 
 All available configuration options
 ===================================

locust/argument_parser.py

@@ -12,18 +12,24 @@
 import sys
 import tempfile
 import textwrap
+from collections import OrderedDict
 from typing import Any, NamedTuple
 from urllib.parse import urlparse
 from uuid import uuid4
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 import configargparse
 import gevent
 import requests
 
 version = locust.__version__
 
 
-DEFAULT_CONFIG_FILES = ["~/.locust.conf", "locust.conf"]
+DEFAULT_CONFIG_FILES = ("~/.locust.conf", "locust.conf", "pyproject.toml")
 
 
 class LocustArgumentParser(configargparse.ArgumentParser):
@@ -63,6 +69,31 @@ def secret_args_included_in_web_ui(self) -> dict[str, configargparse.Action]:
         }
 
 
+class LocustTomlConfigParser(configargparse.TomlConfigParser):
+    def parse(self, stream):
+        try:
+            config = tomllib.loads(stream.read())
+        except Exception as e:
+            raise configargparse.ConfigFileParserException(f"Couldn't parse TOML file: {e}")
+
+        # convert to dict and filter based on section names
+        result = OrderedDict()
+
+        for section in self.sections:
+            data = configargparse.get_toml_section(config, section)
+            if data:
+                for key, value in data.items():
+                    if isinstance(value, list):
+                        result[key] = value
+                    elif value is None:
+                        pass
+                    else:
+                        result[key] = str(value)
+                break
+
+        return result
+
+
 def _is_package(path):
     """
     Is the given path a Python package?
@@ -186,6 +217,12 @@ def exit_handler():
 def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser:
     parser = LocustArgumentParser(
         default_config_files=default_config_files,
+        config_file_parser_class=configargparse.CompositeConfigParser(
+            [
+                LocustTomlConfigParser(["tool.locust"]),
+                configargparse.DefaultConfigFileParser,
+            ]
+        ),
         add_env_var_help=False,
         add_config_file_help=False,
         add_help=add_help,

locust/test/test_parser.py

@@ -38,17 +38,50 @@ def test_skip_log_setup(self):
         opts = self.parser.parse_args(args)
         self.assertEqual(opts.skip_log_setup, True)
 
-    def test_parameter_parsing(self):
-        with NamedTemporaryFile(mode="w") as file:
-            os.environ["LOCUST_LOCUSTFILE"] = "locustfile_from_env"
-            file.write("host host_from_config\nweb-host webhost_from_config")
+    def test_parse_options_from_conf_file(self):
+        with NamedTemporaryFile(mode="w", suffix=".conf") as file:
+            config_data = """\
+            locustfile = ./test_locustfile.py
+            web-host = 127.0.0.1
+            web-port = 45787
+            headless
+            tags = [Critical, Normal]
+            """
+
+            file.write(config_data)
             file.flush()
             parser = get_parser(default_config_files=[file.name])
-            options = parser.parse_args(["-H", "host_from_args"])
-            del os.environ["LOCUST_LOCUSTFILE"]
-        self.assertEqual(options.web_host, "webhost_from_config")
-        self.assertEqual(options.locustfile, "locustfile_from_env")
-        self.assertEqual(options.host, "host_from_args")  # overridden
+            options = parser.parse_args(["-H", "https://example.com"])
+
+        self.assertEqual("./test_locustfile.py", options.locustfile)
+        self.assertEqual("127.0.0.1", options.web_host)
+        self.assertEqual(45787, options.web_port)
+        self.assertTrue(options.headless)
+        self.assertEqual(["Critical", "Normal"], options.tags)
+        self.assertEqual("https://example.com", options.host)
+
+    def test_parse_options_from_toml_file(self):
+        with NamedTemporaryFile(mode="w", suffix=".toml") as file:
+            config_data = """\
+            [tool.locust]
+            locustfile = "./test_locustfile.py"
+            web-host = "127.0.0.1"
+            web-port = 45787
+            headless = true
+            tags = ["Critical", "Normal"]
+            """
+
+            file.write(config_data)
+            file.flush()
+            parser = get_parser(default_config_files=[file.name])
+            options = parser.parse_args(["-H", "https://example.com"])
+
+        self.assertEqual("./test_locustfile.py", options.locustfile)
+        self.assertEqual("127.0.0.1", options.web_host)
+        self.assertEqual(45787, options.web_port)
+        self.assertTrue(options.headless)
+        self.assertEqual(["Critical", "Normal"], options.tags)
+        self.assertEqual("https://example.com", options.host)
 
 
 class TestArgumentParser(LocustTestCase):

pyproject.toml

@@ -17,11 +17,12 @@ dependencies = [
     "pyzmq >=25.0.0",
     "geventhttpclient >=2.0.11",
     "ConfigArgParse >=1.5.5",
+    "tomli >=1.1.0; python_version<'3.11'",
     "psutil >=5.9.1",
     "Flask-Login >=0.6.3",
     "Flask-Cors >=3.0.10",
     "roundrobin >=0.0.2",
-    "pywin32;platform_system=='Windows'",
+    "pywin32; platform_system=='Windows'",
 ]
 classifiers = [
     "Topic :: Software Development :: Testing :: Traffic Generation",