locust 2.20.1.dev28__py3-none-any.whl → 2.20.2__py3-none-any.whl

locust/__init__.py

@@ -15,15 +15,15 @@ from gevent import monkey
 monkey.patch_all()
 
 from ._version import version as __version__
-from .user.sequential_taskset import SequentialTaskSet
-from .user import wait_time
-from .user.task import task, tag, TaskSet
-from .user.users import HttpUser, User
 from .contrib.fasthttp import FastHttpUser
-from .user.wait_time import between, constant, constant_pacing, constant_throughput
-from .shape import LoadTestShape
 from .debug import run_single_user
 from .event import Events
+from .shape import LoadTestShape
+from .user import wait_time
+from .user.sequential_taskset import SequentialTaskSet
+from .user.task import TaskSet, tag, task
+from .user.users import HttpUser, User
+from .user.wait_time import between, constant, constant_pacing, constant_throughput
 
 events = Events()
 
@@ -46,5 +46,5 @@ __all__ = (
 )
 
 # Used for raising a DeprecationWarning if old Locust/HttpLocust is used
-from .util.deprecation import DeprecatedLocustClass as Locust
 from .util.deprecation import DeprecatedHttpLocustClass as HttpLocust
+from .util.deprecation import DeprecatedLocustClass as Locust

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.20.1.dev28'
-__version_tuple__ = version_tuple = (2, 20, 1, 'dev28')
+__version__ = version = '2.20.2'
+__version_tuple__ = version_tuple = (2, 20, 2)

locust/argument_parser.py

@@ -1,11 +1,14 @@
+from __future__ import annotations
+
+import locust
+
 import os
 import platform
 import sys
 import textwrap
-from typing import Dict, List, NamedTuple, Optional, Any
-import configargparse
+from typing import Any, NamedTuple
 
-import locust
+import configargparse
 
 version = locust.__version__
 
@@ -38,11 +41,11 @@ class LocustArgumentParser(configargparse.ArgumentParser):
         return action
 
     @property
-    def args_included_in_web_ui(self) -> Dict[str, configargparse.Action]:
+    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}
 
     @property
-    def secret_args_included_in_web_ui(self) -> Dict[str, configargparse.Action]:
+    def secret_args_included_in_web_ui(self) -> dict[str, configargparse.Action]:
         return {
             a.dest: a
             for a in self._actions
@@ -91,7 +94,7 @@ def find_locustfile(locustfile):
     # Implicit 'return None' if nothing was found
 
 
-def find_locustfiles(locustfiles: List[str], is_directory: bool) -> List[str]:
+def find_locustfiles(locustfiles: list[str], is_directory: bool) -> list[str]:
     """
     Returns a list of relative file paths for the Locustfile Picker. If is_directory is True,
     locustfiles is expected to have a single index which is a directory that will be searched for
@@ -176,7 +179,7 @@ See documentation for more details, including how to set options using a file or
     return parser
 
 
-def parse_locustfile_option(args=None) -> List[str]:
+def parse_locustfile_option(args=None) -> list[str]:
     """
     Construct a command line parser that is only used to parse the -f argument so that we can
     import the test scripts in case any of them adds additional command line arguments to the
@@ -304,6 +307,13 @@ def setup_parser_arguments(parser):
         dest="list_commands",
         help="Show list of possible User classes and exit",
     )
+    parser.add_argument(
+        "--config-users",
+        type=str,
+        nargs="*",
+        help="User configuration as a JSON string or file. A list of arguments or an Array of JSON configuration may be provided",
+        env_var="LOCUST_CONFIG_USERS",
+    )
 
     web_ui_group = parser.add_argument_group("Web UI options")
     web_ui_group.add_argument(
@@ -355,9 +365,16 @@ def setup_parser_arguments(parser):
         dest="web_auth",
         metavar="<username:password>",
         default=None,
-        help="DEPRECATED. See https://github.com/locustio/locust/issues/2517 Turn on Basic Auth for the web interface. Should be supplied in the following format: username:password ",
+        help=configargparse.SUPPRESS,
         env_var="LOCUST_WEB_AUTH",
     )
+    web_ui_group.add_argument(
+        "--web-login",
+        default=False,
+        action="store_true",
+        help="Protects the web interface with a login page. See https://docs.locust.io/en/stable/extending-locust.html#authentication",
+        env_var="LOCUST_WEB_LOGIN",
+    )
     web_ui_group.add_argument(
         "--tls-cert",
         default="",
@@ -429,6 +446,13 @@ def setup_parser_arguments(parser):
         help="How long should the master wait for workers to connect before giving up. Defaults to wait forever",
         env_var="LOCUST_EXPECT_WORKERS_MAX_WAIT",
     )
+    master_group.add_argument(
+        "--enable-rebalancing",
+        action="store_true",
+        default=False,
+        dest="enable_rebalancing",
+        help="Re-distribute users if new workers are added or removed during a test run. Experimental.",
+    )
     master_group.add_argument(
         "--expect-slaves",
         action="store_true",
@@ -611,13 +635,6 @@ Typically ONLY these options (and --locustfile) need to be specified on workers,
         dest="equal_weights",
         help="Use equally distributed task weights, overriding the weights specified in the locustfile.",
     )
-    other_group.add_argument(
-        "--enable-rebalancing",
-        action="store_true",
-        default=False,
-        dest="enable_rebalancing",
-        help="Allow to automatically rebalance users if new workers are added or removed during a test run.",
-    )
 
     user_classes_group = parser.add_argument_group("User classes")
     user_classes_group.add_argument(
@@ -660,10 +677,10 @@ class UIExtraArgOptions(NamedTuple):
     default_value: str
     is_secret: bool
     help_text: str
-    choices: Optional[List[str]] = None
+    choices: list[str] | None = None
 
 
-def ui_extra_args_dict(args=None) -> Dict[str, Dict[str, Any]]:
+def ui_extra_args_dict(args=None) -> dict[str, dict[str, Any]]:
     """Get all the UI visible arguments"""
     locust_args = default_args_dict()
 
@@ -684,7 +701,7 @@ def ui_extra_args_dict(args=None) -> Dict[str, Dict[str, Any]]:
     return extra_args
 
 
-def locustfile_is_directory(locustfiles: List[str]) -> bool:
+def locustfile_is_directory(locustfiles: list[str]) -> bool:
     """
     If a user passes in a locustfile without a file extension and there is a directory with the same name,
     this function defaults to using the file and will raise a warning.

locust/clients.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
+
 import re
 import time
 from contextlib import contextmanager
-from typing import Generator, Optional
+from typing import Generator
 from urllib.parse import urlparse, urlunparse
 
 import requests
@@ -48,7 +49,7 @@ class HttpSession(requests.Session):
                            and then mark it as successful even if the response code was not (i.e 500 or 404).
     """
 
-    def __init__(self, base_url, request_event, user, *args, pool_manager: Optional[PoolManager] = None, **kwargs):
+    def __init__(self, base_url, request_event, user, *args, pool_manager: PoolManager | None = None, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.base_url = base_url
@@ -57,7 +58,7 @@ class HttpSession(requests.Session):
 
         # User can group name, or use the group context manager to gather performance statistics under a specific name
         # This is an alternative to passing in the "name" parameter to the requests function
-        self.request_name: Optional[str] = None
+        self.request_name: str | None = None
 
         # Check for basic authentication
         parsed_url = urlparse(self.base_url)
@@ -301,7 +302,7 @@ class ResponseContextManager(LocustResponse):
 
 
 class LocustHttpAdapter(HTTPAdapter):
-    def __init__(self, pool_manager: Optional[PoolManager], *args, **kwargs):
+    def __init__(self, pool_manager: PoolManager | None, *args, **kwargs):
         self.poolmanager = pool_manager
         super().__init__(*args, **kwargs)
 
@@ -323,5 +324,5 @@ def _failure(self):
     )
 
 
-Response.success = _success
-Response.failure = _failure
+Response.success = _success  # type: ignore[attr-defined]
+Response.failure = _failure  # type: ignore[attr-defined]

locust/contrib/fasthttp.py

@@ -1,35 +1,35 @@
 from __future__ import annotations
-import re
-import socket
+
+from locust.env import Environment
+from locust.exception import CatchResponseError, LocustError, ResponseError
+from locust.user import User
+from locust.util.deprecation import DeprecatedFastHttpLocustClass as FastHttpLocust
+
 import json
 import json as unshadowed_json  # some methods take a named parameter called json
+import re
+import socket
+import time
+import traceback
 from base64 import b64encode
 from contextlib import contextmanager
+from http.cookiejar import CookieJar
 from json.decoder import JSONDecodeError
-from urllib.parse import urlparse, urlunparse
 from ssl import SSLError
-import time
-import traceback
-from typing import Callable, Optional, Tuple, Dict, Any, Generator, cast
-
-from http.cookiejar import CookieJar
+from typing import Any, Callable, Generator, cast
+from urllib.parse import urlparse, urlunparse
 
 import gevent
+from charset_normalizer import detect
 from gevent.timeout import Timeout
 from geventhttpclient._parser import HTTPParseError
 from geventhttpclient.client import HTTPClientPool
-from geventhttpclient.useragent import UserAgent, CompatRequest, CompatResponse, ConnectionError
-from geventhttpclient.response import HTTPConnectionClosed, HTTPSocketPoolResponse
 from geventhttpclient.header import Headers
+from geventhttpclient.response import HTTPConnectionClosed, HTTPSocketPoolResponse
+from geventhttpclient.useragent import CompatRequest, CompatResponse, ConnectionError, UserAgent
 
 # borrow requests's content-type header parsing
 from requests.utils import get_encoding_from_headers
-from charset_normalizer import detect
-
-from locust.user import User
-from locust.exception import LocustError, CatchResponseError, ResponseError
-from locust.env import Environment
-from locust.util.deprecation import DeprecatedFastHttpLocustClass as FastHttpLocust
 
 # Monkey patch geventhttpclient.useragent.CompatRequest so that Cookiejar works with Python >= 3.3
 # More info: https://github.com/requests/requests/pull/871
@@ -79,10 +79,10 @@ class FastHttpSession:
         self,
         environment: Environment,
         base_url: str,
-        user: Optional[User],
+        user: User | None,
         insecure=True,
-        client_pool: Optional[HTTPClientPool] = None,
-        ssl_context_factory: Optional[Callable] = None,
+        client_pool: HTTPClientPool | None = None,
+        ssl_context_factory: Callable | None = None,
         **kwargs,
     ):
         self.environment = environment
@@ -313,7 +313,7 @@ class FastHttpUser(User):
     insecure: bool = True
     """Parameter passed to FastHttpSession. Default True, meaning no SSL verification."""
 
-    default_headers: Optional[dict] = None
+    default_headers: dict | None = None
     """Parameter passed to FastHttpSession. Adds the listed headers to every request."""
 
     concurrency: int = 10
@@ -321,10 +321,10 @@ class FastHttpUser(User):
     Note that setting this value has no effect when custom client_pool was given, and you need to spawn a your own gevent pool
     to use it (as Users only have one greenlet). See test_fasthttp.py / test_client_pool_concurrency for an example."""
 
-    client_pool: Optional[HTTPClientPool] = None
+    client_pool: HTTPClientPool | None = None
     """HTTP client pool to use. If not given, a new pool is created per single user."""
 
-    ssl_context_factory: Optional[Callable] = None
+    ssl_context_factory: Callable | None = None
     """A callable that return a SSLContext for overriding the default context created by the FastHttpSession."""
 
     abstract = True
@@ -360,7 +360,7 @@ class FastHttpUser(User):
 
     @contextmanager
     def rest(
-        self, method, url, headers: Optional[dict] = None, **kwargs
+        self, method, url, headers: dict | None = None, **kwargs
     ) -> Generator[RestResponseContextManager, None, None]:
         """
         A wrapper for self.client.request that:
@@ -423,36 +423,36 @@ class FastHttpUser(User):
 
 
 class FastRequest(CompatRequest):
-    payload: Optional[str] = None
+    payload: str | None = None
 
     @property
-    def body(self) -> Optional[str]:
+    def body(self) -> str | None:
         return self.payload
 
 
 class FastResponse(CompatResponse):
-    headers: Optional[Headers] = None
+    headers: Headers | None = None
     """Dict like object containing the response headers"""
 
-    _response: Optional[HTTPSocketPoolResponse] = None
+    _response: HTTPSocketPoolResponse | None = None
 
-    encoding: Optional[str] = None
+    encoding: str | None = None
     """In some cases setting the encoding explicitly is needed. If so, do it before calling .text"""
 
-    request: Optional[FastRequest] = None
+    request: FastRequest | None = None
 
     def __init__(
         self,
         ghc_response: HTTPSocketPoolResponse,
-        request: Optional[FastRequest] = None,
-        sent_request: Optional[str] = None,
+        request: FastRequest | None = None,
+        sent_request: str | None = None,
     ):
         super().__init__(ghc_response, request, sent_request)
 
         self.request = request
 
     @property
-    def text(self) -> Optional[str]:
+    def text(self) -> str | None:
         """
         Returns the text content of the response as a decoded string
         """
@@ -472,7 +472,7 @@ class FastResponse(CompatResponse):
         return str(self.content, str(self.encoding), errors="replace")
 
     @property
-    def url(self) -> Optional[str]:
+    def url(self) -> str | None:
         """
         Get "response" URL, which is the same as the request URL. This is a small deviation from HttpSession, which gets the final (possibly redirected) URL.
         """
@@ -500,6 +500,11 @@ class FastResponse(CompatResponse):
         """
         return self._response.get_code() if self._response is not None else 0
 
+    @property
+    def ok(self):
+        """Returns True if :attr:`status_code` is less than 400, False if not."""
+        return self.status_code < 400
+
     def _content(self):
         if self.headers is None:
             return None
@@ -522,11 +527,11 @@ class ErrorResponse:
     that doesn't have a real Response object attached. E.g. a socket error or similar
     """
 
-    headers: Optional[Headers] = None
+    headers: Headers | None = None
     content = None
     status_code = 0
-    error: Optional[Exception] = None
-    text: Optional[str] = None
+    error: Exception | None = None
+    text: str | None = None
     request: CompatRequest
 
     def __init__(self, url: str, request: CompatRequest):
@@ -542,7 +547,7 @@ class LocustUserAgent(UserAgent):
     request_type = FastRequest
     valid_response_codes = frozenset([200, 201, 202, 203, 204, 205, 206, 207, 208, 226, 301, 302, 303, 304, 307])
 
-    def __init__(self, client_pool: Optional[HTTPClientPool] = None, **kwargs):
+    def __init__(self, client_pool: HTTPClientPool | None = None, **kwargs):
         super().__init__(**kwargs)
 
         if client_pool is not None:

locust/debug.py

@@ -1,13 +1,19 @@
-from datetime import datetime, timezone
-import os
-import inspect
+from __future__ import annotations
+
 import locust
 import locust.log
-from locust import User, argument_parser
-from typing import Type, Optional
+from locust import argument_parser
 from locust.env import Environment
 from locust.exception import CatchResponseError, RescheduleTask
 
+import inspect
+import os
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from locust import User
+
 
 def _print_t(s):
     """
@@ -94,16 +100,16 @@ class PrintListener:
         print()
 
 
-_env: Optional[Environment] = None  # minimal Environment for debugging
+_env: Environment | None = None  # minimal Environment for debugging
 
 
 def run_single_user(
-    user_class: Type[User],
+    user_class: type[User],
     include_length=False,
     include_time=False,
     include_context=False,
     include_payload=False,
-    loglevel: Optional[str] = "WARNING",
+    loglevel: str | None = "WARNING",
 ):
     """
     Runs a single User. Useful when you want to run a debugger.

locust/dispatch.py

@@ -1,19 +1,19 @@
-from collections import defaultdict
+from __future__ import annotations
+
 import contextlib
 import itertools
 import math
 import time
+from collections import defaultdict
 from collections.abc import Iterator
 from operator import attrgetter
-from typing import Dict, Generator, List, TYPE_CHECKING, Optional, Tuple, Type, Set
+from typing import TYPE_CHECKING, Generator
 
 import gevent
-
 from roundrobin import smooth
 
-from locust import User
-
 if TYPE_CHECKING:
+    from locust import User
     from locust.runners import WorkerNode
 
 
@@ -49,7 +49,7 @@ class UsersDispatcher(Iterator):
             from 10 to 100.
     """
 
-    def __init__(self, worker_nodes: List["WorkerNode"], user_classes: List[Type[User]]):
+    def __init__(self, worker_nodes: list[WorkerNode], user_classes: list[type[User]]):
         """
         :param worker_nodes: List of worker nodes
         :param user_classes: The user classes
@@ -79,16 +79,16 @@ class UsersDispatcher(Iterator):
 
         self._current_user_count = self.get_current_user_count()
 
-        self._dispatcher_generator: Generator[Dict[str, Dict[str, int]], None, None] = None
+        self._dispatcher_generator: Generator[dict[str, dict[str, int]], None, None] = None
 
         self._user_generator = self._user_gen()
 
         self._worker_node_generator = itertools.cycle(self._worker_nodes)
 
         # To keep track of how long it takes for each dispatch iteration to compute
-        self._dispatch_iteration_durations: List[float] = []
+        self._dispatch_iteration_durations: list[float] = []
 
-        self._active_users: List[Tuple[WorkerNode, str]] = []
+        self._active_users: list[tuple[WorkerNode, str]] = []
 
         # TODO: Test that attribute is set when dispatching and unset when done dispatching
         self._dispatch_in_progress = False
@@ -108,10 +108,10 @@ class UsersDispatcher(Iterator):
         return self._dispatch_in_progress
 
     @property
-    def dispatch_iteration_durations(self) -> List[float]:
+    def dispatch_iteration_durations(self) -> list[float]:
         return self._dispatch_iteration_durations
 
-    def __next__(self) -> Dict[str, Dict[str, int]]:
+    def __next__(self) -> dict[str, dict[str, int]]:
         users_on_workers = next(self._dispatcher_generator)
         # TODO: Is this necessary to copy the users_on_workers if we know
         #       it won't be mutated by external code?
@@ -122,7 +122,7 @@ class UsersDispatcher(Iterator):
         worker_nodes_by_id = sorted(self._worker_nodes, key=lambda w: w.id)
 
         # Give every worker an index indicating how many workers came before it on that host
-        workers_per_host = defaultdict(lambda: 0)
+        workers_per_host = defaultdict(int)
         for worker_node in worker_nodes_by_id:
             host = worker_node.id.split("_")[0]
             worker_node._index_within_host = workers_per_host[host]
@@ -131,7 +131,7 @@ class UsersDispatcher(Iterator):
         # Sort again, first by index within host, to ensure Users get started evenly across hosts
         self._worker_nodes = sorted(self._worker_nodes, key=lambda worker: (worker._index_within_host, worker.id))
 
-    def _dispatcher(self) -> Generator[Dict[str, Dict[str, int]], None, None]:
+    def _dispatcher(self) -> Generator[dict[str, dict[str, int]], None, None]:
         self._dispatch_in_progress = True
 
         if self._rebalance:
@@ -164,7 +164,7 @@ class UsersDispatcher(Iterator):
 
         self._dispatch_in_progress = False
 
-    def new_dispatch(self, target_user_count: int, spawn_rate: float, user_classes: Optional[List] = None) -> None:
+    def new_dispatch(self, target_user_count: int, spawn_rate: float, user_classes: list | None = None) -> None:
         """
         Initialize a new dispatch cycle.
 
@@ -194,7 +194,7 @@ class UsersDispatcher(Iterator):
 
         self._dispatch_iteration_durations.clear()
 
-    def add_worker(self, worker_node: "WorkerNode") -> None:
+    def add_worker(self, worker_node: WorkerNode) -> None:
         """
         This method is to be called when a new worker connects to the master. When
         a new worker is added, the users dispatcher will flag that a rebalance is required
@@ -207,7 +207,7 @@ class UsersDispatcher(Iterator):
         self._sort_workers()
         self._prepare_rebalance()
 
-    def remove_worker(self, worker_node: "WorkerNode") -> None:
+    def remove_worker(self, worker_node: WorkerNode) -> None:
         """
         This method is similar to the above `add_worker`. When a worker disconnects
         (because of e.g. network failure, worker failure, etc.), this method will ensure that the next
@@ -268,7 +268,7 @@ class UsersDispatcher(Iterator):
         sleep_duration = max(0.0, self._wait_between_dispatch - delta)
         gevent.sleep(sleep_duration)
 
-    def _add_users_on_workers(self) -> Dict[str, Dict[str, int]]:
+    def _add_users_on_workers(self) -> dict[str, dict[str, int]]:
         """Add users on the workers until the target number of users is reached for the current dispatch iteration
 
         :return: The users that we want to run on the workers
@@ -290,7 +290,7 @@ class UsersDispatcher(Iterator):
 
         return self._users_on_workers
 
-    def _remove_users_from_workers(self) -> Dict[str, Dict[str, int]]:
+    def _remove_users_from_workers(self) -> dict[str, dict[str, int]]:
         """Remove users from the workers until the target number of users is reached for the current dispatch iteration
 
         :return: The users that we want to run on the workers
@@ -318,8 +318,8 @@ class UsersDispatcher(Iterator):
 
     def _distribute_users(
         self, target_user_count: int
-    ) -> Tuple[
-        Dict[str, Dict[str, int]], Generator[Optional[str], None, None], itertools.cycle, List[Tuple["WorkerNode", str]]
+    ) -> tuple[
+        dict[str, dict[str, int]], Generator[str | None, None, None], itertools.cycle, list[tuple[WorkerNode, str]]
     ]:
         """
         This function might take some time to complete if the `target_user_count` is a big number. A big number
@@ -349,7 +349,7 @@ class UsersDispatcher(Iterator):
 
         return users_on_workers, user_gen, worker_gen, active_users
 
-    def _user_gen(self) -> Generator[Optional[str], None, None]:
+    def _user_gen(self) -> Generator[str | None, None, None]:
         """
         This method generates users according to their weights using
         a smooth weighted round-robin algorithm implemented by https://github.com/linnik/roundrobin.
@@ -361,7 +361,7 @@ class UsersDispatcher(Iterator):
         less accurate during ramp-up/down.
         """
 
-        def infinite_cycle_gen(users: List[Tuple[Type[User], int]]) -> itertools.cycle:
+        def infinite_cycle_gen(users: list[tuple[type[User], int]]) -> itertools.cycle:
             if not users:
                 return itertools.cycle([None])
 
@@ -401,9 +401,9 @@ class UsersDispatcher(Iterator):
             if self._try_dispatch_fixed:
                 self._try_dispatch_fixed = False
                 current_fixed_users_count = {u: self._get_user_current_count(u) for u in fixed_users}
-                spawned_classes: Set[str] = set()
+                spawned_classes: set[str] = set()
                 while len(spawned_classes) != len(fixed_users):
-                    user_name: Optional[str] = next(cycle_fixed_gen)
+                    user_name: str | None = next(cycle_fixed_gen)
                     if not user_name:
                         break
 
@@ -422,7 +422,7 @@ class UsersDispatcher(Iterator):
             yield next(cycle_weighted_gen)
 
     @staticmethod
-    def _fast_users_on_workers_copy(users_on_workers: Dict[str, Dict[str, int]]) -> Dict[str, Dict[str, int]]:
+    def _fast_users_on_workers_copy(users_on_workers: dict[str, dict[str, int]]) -> dict[str, dict[str, int]]:
         """deepcopy is too slow, so we use this custom copy function.
 
         The implementation was profiled and compared to other implementations such as dict-comprehensions