Add more type hints

locust/contrib/fasthttp.py

@@ -54,66 +54,6 @@ def _construct_basic_auth_str(username, password):
     return "Basic " + b64encode(b":".join((username, password))).strip().decode("ascii")
 
 
-class FastHttpUser(User):
-    """
-    FastHttpUser uses a different HTTP client (geventhttpclient) compared to HttpUser (python-requests).
-    It's significantly faster, but not as capable.
-
-    The behaviour of this user is defined by it's tasks. Tasks can be declared either directly on the
-    class by using the :py:func:`@task decorator <locust.task>` on the methods, or by setting
-    the :py:attr:`tasks attribute <locust.User.tasks>`.
-
-    This class creates a *client* attribute on instantiation which is an HTTP client with support
-    for keeping a user session between requests.
-    """
-
-    client = None
-    """
-    Instance of HttpSession that is created upon instantiation of User.
-    The client support cookies, and therefore keeps the session between HTTP requests.
-    """
-
-    # Below are various UserAgent settings. Change these in your subclass to alter FastHttpUser's behaviour.
-    # It needs to be done before FastHttpUser is instantiated, changing them later will have no effect
-
-    network_timeout: float = 60.0
-    """Parameter passed to FastHttpSession"""
-
-    connection_timeout: float = 60.0
-    """Parameter passed to FastHttpSession"""
-
-    max_redirects: int = 5
-    """Parameter passed to FastHttpSession. Default 5, meaning 4 redirects."""
-
-    max_retries: int = 1
-    """Parameter passed to FastHttpSession. Default 1, meaning zero retries."""
-
-    insecure: bool = True
-    """Parameter passed to FastHttpSession. Default True, meaning no SSL verification."""
-
-    abstract = True
-    """Dont register this as a User class that can be run by itself"""
-
-    def __init__(self, environment):
-        super().__init__(environment)
-        if self.host is None:
-            raise LocustError(
-                "You must specify the base host. Either in the host attribute in the User class, or on the command line using the --host option."
-            )
-        if not re.match(r"^https?://[^/]+", self.host, re.I):
-            raise LocustError("Invalid host (`%s`), must be a valid base URL. E.g. http://example.com" % self.host)
-
-        self.client = FastHttpSession(
-            self.environment,
-            base_url=self.host,
-            network_timeout=self.network_timeout,
-            connection_timeout=self.connection_timeout,
-            max_redirects=self.max_redirects,
-            max_retries=self.max_retries,
-            insecure=self.insecure,
-        )
-
-
 def insecure_ssl_context_factory():
     context = gevent.ssl.create_default_context()
     context.check_hostname = False
@@ -325,6 +265,66 @@ def put(self, path, data=None, **kwargs):
         return self.request("PUT", path, data=data, **kwargs)
 
 
+class FastHttpUser(User):
+    """
+    FastHttpUser uses a different HTTP client (geventhttpclient) compared to HttpUser (python-requests).
+    It's significantly faster, but not as capable.
+
+    The behaviour of this user is defined by it's tasks. Tasks can be declared either directly on the
+    class by using the :py:func:`@task decorator <locust.task>` on the methods, or by setting
+    the :py:attr:`tasks attribute <locust.User.tasks>`.
+
+    This class creates a *client* attribute on instantiation which is an HTTP client with support
+    for keeping a user session between requests.
+    """
+
+    client: FastHttpSession = None
+    """
+    Instance of HttpSession that is created upon instantiation of User.
+    The client support cookies, and therefore keeps the session between HTTP requests.
+    """
+
+    # Below are various UserAgent settings. Change these in your subclass to alter FastHttpUser's behaviour.
+    # It needs to be done before FastHttpUser is instantiated, changing them later will have no effect
+
+    network_timeout: float = 60.0
+    """Parameter passed to FastHttpSession"""
+
+    connection_timeout: float = 60.0
+    """Parameter passed to FastHttpSession"""
+
+    max_redirects: int = 5
+    """Parameter passed to FastHttpSession. Default 5, meaning 4 redirects."""
+
+    max_retries: int = 1
+    """Parameter passed to FastHttpSession. Default 1, meaning zero retries."""
+
+    insecure: bool = True
+    """Parameter passed to FastHttpSession. Default True, meaning no SSL verification."""
+
+    abstract = True
+    """Dont register this as a User class that can be run by itself"""
+
+    def __init__(self, environment):
+        super().__init__(environment)
+        if self.host is None:
+            raise LocustError(
+                "You must specify the base host. Either in the host attribute in the User class, or on the command line using the --host option."
+            )
+        if not re.match(r"^https?://[^/]+", self.host, re.I):
+            raise LocustError("Invalid host (`%s`), must be a valid base URL. E.g. http://example.com" % self.host)
+
+        self.client = FastHttpSession(
+            self.environment,
+            base_url=self.host,
+            network_timeout=self.network_timeout,
+            connection_timeout=self.connection_timeout,
+            max_redirects=self.max_redirects,
+            max_retries=self.max_retries,
+            insecure=self.insecure,
+        )
+
+
 class FastResponse(CompatResponse):
     headers = None
     """Dict like object containing the response headers"""

locust/env.py

@@ -1,23 +1,25 @@
 from .event import Events
 from .exception import RunnerAlreadyExistsError
 from .stats import RequestStats
-from .runners import LocalRunner, MasterRunner, Runner, WorkerRunner
+from .runners import Runner, LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
+from .user import User
 from .user.task import filter_tasks_by_tags
 from .shape import LoadTestShape
+from typing import List
 
 
 class Environment:
-    events = None
+    events: Events = None
     """
     Event hooks used by Locust internally, as well as to extend Locust's functionality
     See :ref:`events` for available events.
     """
 
-    user_classes = []
+    user_classes: List[User] = []
     """User classes that the runner will run"""
 
-    shape_class = None
+    shape_class: LoadTestShape = None
     """A shape class to control the shape of the load test"""
 
     tags = None
@@ -26,16 +28,16 @@ class Environment:
     exclude_tags = None
     """If set, only tasks that aren't tagged by tags in this list will be executed"""
 
-    stats = None
+    stats: RequestStats = None
     """Reference to RequestStats instance"""
 
     runner: Runner = None
     """Reference to the :class:`Runner <locust.runners.Runner>` instance"""
 
-    web_ui = None
+    web_ui: WebUI = None
     """Reference to the WebUI instance"""
 
-    host = None
+    host: str = None
     """Base URL of the target system"""
 
     reset_stats = False

locust/shape.py

@@ -1,4 +1,5 @@
 import time
+from typing import Optional, Tuple
 
 
 class LoadTestShape(object):
@@ -22,7 +23,7 @@ def get_run_time(self):
         """
         return time.monotonic() - self.start_time
 
-    def tick(self):
+    def tick(self) -> Optional[Tuple[int, float]]:
         """
         Returns a tuple with 2 elements to control the running load test:
 

locust/test/test_locust_class.py

@@ -698,22 +698,6 @@ def t1(l):
         self.assertEqual(1, self.runner.stats.get("new name!", "GET").num_requests)
         self.assertEqual(0, self.runner.stats.get("/ultra_fast", "GET").num_requests)
 
-    def test_locust_client_error(self):
-        class MyTaskSet(TaskSet):
-            @task
-            def t1(self):
-                self.client.get("/")
-                self.interrupt()
-
-        class MyUser(User):
-            host = "http://127.0.0.1:%i" % self.port
-            tasks = [MyTaskSet]
-
-        my_locust = MyUser(self.environment)
-        self.assertRaises(LocustError, lambda: my_locust.client.get("/"))
-        my_taskset = MyTaskSet(my_locust)
-        self.assertRaises(LocustError, lambda: my_taskset.client.get("/"))
-
     def test_redirect_url_original_path_as_name(self):
         class MyUser(HttpUser):
             host = "http://127.0.0.1:%i" % self.port

locust/test/test_main.py

@@ -52,7 +52,9 @@ def test_load_locust_file_from_absolute_path(self):
 
     def test_load_locust_file_from_relative_path(self):
         with mock_locustfile() as mocked:
-            docstring, user_classes, shape_class = main.load_locustfile(os.path.join("./locust/test/", mocked.filename))
+            docstring, user_classes, shape_class = main.load_locustfile(
+                os.path.join(os.path.relpath(mocked.directory, os.getcwd()), mocked.filename)
+            )
 
     def test_load_locust_file_with_a_dot_in_filename(self):
         with mock_locustfile(filename_prefix="mocked.locust.file") as mocked:

locust/user/task.py

@@ -3,6 +3,7 @@
 import sys
 import traceback
 from time import time
+from typing import Any, Callable, List, Union
 
 import gevent
 from gevent import GreenletExit
@@ -185,7 +186,7 @@ class TaskSet(object, metaclass=TaskSetMeta):
     will then continue in the first TaskSet).
     """
 
-    tasks = []
+    tasks: List[Union["TaskSet", Callable]] = []
     """
     Collection of python callables and/or TaskSet classes that the User(s) will run.
 

locust/user/users.py

@@ -1,9 +1,11 @@
+from typing import Any, Callable, List, TypeVar, Union
 from gevent import GreenletExit, greenlet
 from gevent.pool import Group
 from locust.clients import HttpSession
 from locust.exception import LocustError, StopUser
 from locust.util import deprecation
 from .task import (
+    TaskSet,
     DefaultTaskSet,
     get_tasks_from_base_classes,
     LOCUST_STATE_RUNNING,
@@ -12,16 +14,6 @@
 )
 
 
-class NoClientWarningRaiser(object):
-    """
-    The purpose of this class is to emit a sensible error message for old test scripts that
-    inherit from User, and expects there to be an HTTP client under the client attribute.
-    """
-
-    def __getattr__(self, _):
-        raise LocustError("No client instantiated. Did you intend to inherit from HttpUser?")
-
-
 class UserMeta(type):
     """
     Meta class for the main User class. It's used to allow User classes to specify task execution
@@ -55,7 +47,7 @@ class by using the :py:func:`@task decorator <locust.task>` on methods, or by se
     :py:class:`HttpUser <locust.HttpUser>` class.
     """
 
-    host = None
+    host: str = None
     """Base hostname to swarm. i.e: http://127.0.0.1:1234"""
 
     min_wait = None
@@ -85,7 +77,7 @@ class MyUser(User):
     Method that returns the time between the execution of locust tasks in milliseconds
     """
 
-    tasks = []
+    tasks: List[Union[TaskSet, Callable]] = []
     """
     Collection of python callables and/or TaskSet classes that the Locust user(s) will run.
 
@@ -109,7 +101,7 @@ class ForumPage(TaskSet):
     environment = None
     """A reference to the :py:attr:`environment <locust.Environment>` in which this locust is running"""
 
-    client = NoClientWarningRaiser()
+    client = None
     _state = None
     _greenlet: greenlet.Greenlet = None
     _group: Group
@@ -208,7 +200,7 @@ class by using the :py:func:`@task decorator <locust.task>` on methods, or by se
     abstract = True
     """If abstract is True, the class is meant to be subclassed, and users will not choose this locust during a test"""
 
-    client = None
+    client: HttpSession = None
     """
     Instance of HttpSession that is created upon instantiation of Locust.
     The client supports cookies, and therefore keeps the session between HTTP requests.