From bf0aacb7db03e46e2b8782605f327afb9aed1c91 Mon Sep 17 00:00:00 2001
From: David Baumgold <david@davidbaumgold.com>
Date: Sun, 15 Oct 2023 16:50:46 +0200
Subject: [PATCH] Support argcomplete as an optional dependency

---
 locust/argument_parser.py | 10 ++++++++++
 pyproject.toml            |  3 +++
 2 files changed, 13 insertions(+)

diff --git a/locust/argument_parser.py b/locust/argument_parser.py
index f6d524a77d..5fe9ba6373 100644
--- a/locust/argument_parser.py
+++ b/locust/argument_parser.py
@@ -5,6 +5,14 @@
 from typing import Dict, List, NamedTuple, Optional, Any
 import configargparse
 
+try:
+    from argcomplete import autocomplete
+except ImportError:
+
+    def autocomplete(parser):
+        return None
+
+
 import locust
 
 version = locust.__version__
@@ -606,6 +614,7 @@ def get_parser(default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParse
 
 def parse_options(args=None) -> configargparse.Namespace:
     parser = get_parser()
+    autocomplete(parser)
     parsed_opts = parser.parse_args(args=args)
     if parsed_opts.stats_history_enabled and (parsed_opts.csv_prefix is None):
         parser.error("'--csv-full-history' requires '--csv'.")
@@ -633,6 +642,7 @@ def ui_extra_args_dict(args=None) -> Dict[str, Dict[str, Any]]:
     locust_args = default_args_dict()
 
     parser = get_parser()
+    autocomplete(parser)
     all_args = vars(parser.parse_args(args))
 
     extra_args = {
diff --git a/pyproject.toml b/pyproject.toml
index 6427931ce3..88625abde7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,6 +42,9 @@ classifiers = [
     "Topic :: System :: Distributed Computing",
 ]
 
+[project.optional-dependencies]
+completion = ["argcomplete"]
+
 [project.urls]
 Homepage = "https://locust.io/"
 Documentation = "https://docs.locust.io/"