From f305cd2312d5ffb3839b190be385157630366fc0 Mon Sep 17 00:00:00 2001
From: Lars Holmberg <Lars.Holmberg@svenskaspel.se>
Date: Mon, 18 Mar 2024 09:23:57 +0100
Subject: [PATCH] fix toml handling

---
 har2locust/argument_parser.py | 38 +++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/har2locust/argument_parser.py b/har2locust/argument_parser.py
index fe8973e..27b7763 100644
--- a/har2locust/argument_parser.py
+++ b/har2locust/argument_parser.py
@@ -1,12 +1,44 @@
+import sys
 from argparse import RawDescriptionHelpFormatter
+from collections import OrderedDict
 
 import configargparse
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 from ._version import version
 
 DEFAULT_CONFIG_FILES = ["~/.har2locust.conf", "har2locust.conf", "pyproject.toml"]
 
 
+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 get_parser() -> configargparse.ArgumentParser:
     parser = configargparse.ArgumentParser(
         epilog="""Simplest usage:
@@ -24,6 +56,12 @@ def get_parser() -> configargparse.ArgumentParser:
         add_config_file_help=False,
         formatter_class=RawDescriptionHelpFormatter,
         default_config_files=DEFAULT_CONFIG_FILES,
+        config_file_parser_class=configargparse.CompositeConfigParser(
+            [
+                LocustTomlConfigParser(["tool.har2locust"]),
+                configargparse.DefaultConfigFileParser,
+            ]
+        ),
     )
     parser.add_argument(
         "input",