From 7694c7b8e9f8e5d08b88c2ed83a200a4b4f54564 Mon Sep 17 00:00:00 2001
From: Lars Holmberg <Lars.Holmberg@svenskaspel.se>
Date: Tue, 19 Mar 2024 20:00:04 +0100
Subject: [PATCH] Fix support for pyproject.toml (code copied over from locust,
 so it should behave the same way as it does)

---
 locust_swarm/swarm.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/locust_swarm/swarm.py b/locust_swarm/swarm.py
index ede815f..dfcd531 100644
--- a/locust_swarm/swarm.py
+++ b/locust_swarm/swarm.py
@@ -12,6 +12,7 @@
 import subprocess
 import sys
 import time
+from collections import OrderedDict
 from datetime import datetime, timezone
 
 import configargparse
@@ -20,6 +21,37 @@
 
 from locust_swarm._version import version
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+
+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
+
+
 logging.basicConfig(
     format="%(asctime)s,%(msecs)d %(levelname)-4s [%(filename)s:%(lineno)d] %(message)s",
     datefmt="%Y-%m-%d:%H:%M:%S",
@@ -38,6 +70,10 @@
     ],
     auto_env_var_prefix="LOCUST_",
     formatter_class=configargparse.RawDescriptionHelpFormatter,
+    config_file_parser_class=configargparse.CompositeConfigParser([
+        LocustTomlConfigParser(["tool.locust"]),
+        configargparse.DefaultConfigFileParser,
+    ]),
     description="""A tool for automating distributed locust runs using ssh.
 
 Example: swarm -f test.py --loadgen-list loadgen1.domain.com,loadgen2.domain.com --users 50""",