Merge pull request #2106 from mgor/feature/locust_typing_round_2

additional typing improvements

locust/dispatch.py

@@ -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
@@ -397,7 +397,7 @@ def infinite_cycle_gen(users: List[Tuple[Type[User], int]]) -> itertools.cycle:
                 current_fixed_users_count = {u: self._get_user_current_count(u) for u in fixed_users}
                 spawned_classes: Set[str] = set()
                 while len(spawned_classes) != len(fixed_users):
-                    user_name = next(cycle_fixed_gen)
+                    user_name: Optional[str] = next(cycle_fixed_gen)
                     if not user_name:
                         break
 

locust/env.py

@@ -5,19 +5,19 @@
     List,
     Type,
     TypeVar,
-    Union,
     Optional,
+    Union,
 )
 
 from configargparse import Namespace
 
 from .event import Events
 from .exception import RunnerAlreadyExistsError
-from .stats import RequestStats
+from .stats import RequestStats, StatsCSV
 from .runners import Runner, LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
 from .user import User
-from .user.task import TaskSet, filter_tasks_by_tags
+from .user.task import filter_tasks_by_tags, TaskSet, TaskHolder
 from .shape import LoadTestShape
 
 
@@ -28,17 +28,17 @@ class Environment:
     def __init__(
         self,
         *,
-        user_classes: Union[List[Type[User]], None] = None,
-        shape_class: Union[LoadTestShape, None] = None,
-        tags: Union[List[str], None] = None,
+        user_classes: Optional[List[Type[User]]] = None,
+        shape_class: Optional[LoadTestShape] = None,
+        tags: Optional[List[str]] = None,
         locustfile: str = None,
-        exclude_tags=None,
+        exclude_tags: Optional[List[str]] = None,
         events: Events = None,
         host: str = None,
         reset_stats=False,
-        stop_timeout: Union[float, None] = None,
+        stop_timeout: Optional[float] = None,
         catch_exceptions=True,
-        parsed_options: Namespace = None,
+        parsed_options: Optional[Namespace] = None,
     ):
 
         self.runner: Optional[Runner] = None
@@ -141,7 +141,7 @@ def create_master_runner(self, master_bind_host="*", master_bind_port=5557) -> M
             master_bind_port=master_bind_port,
         )
 
-    def create_worker_runner(self, master_host, master_port) -> WorkerRunner:
+    def create_worker_runner(self, master_host: str, master_port: int) -> WorkerRunner:
         """
         Create a :class:`WorkerRunner <locust.runners.WorkerRunner>` instance for this Environment
 
@@ -161,12 +161,12 @@ def create_web_ui(
         self,
         host="",
         port=8089,
-        auth_credentials=None,
-        tls_cert=None,
-        tls_key=None,
-        stats_csv_writer=None,
+        auth_credentials: Optional[str] = None,
+        tls_cert: Optional[str] = None,
+        tls_key: Optional[str] = None,
+        stats_csv_writer: Optional[StatsCSV] = None,
         delayed_start=False,
-    ):
+    ) -> WebUI:
         """
         Creates a :class:`WebUI <locust.web.WebUI>` instance for this Environment and start running the web server
 
@@ -194,7 +194,7 @@ def create_web_ui(
         )
         return self.web_ui
 
-    def _filter_tasks_by_tags(self):
+    def _filter_tasks_by_tags(self) -> None:
         """
         Filter the tasks on all the user_classes recursively, according to the tags and
         exclude_tags attributes
@@ -220,7 +220,7 @@ def _filter_tasks_by_tags(self):
         for user_class in self.user_classes:
             filter_tasks_by_tags(user_class, tags, exclude_tags)
 
-    def _remove_user_classes_with_weight_zero(self):
+    def _remove_user_classes_with_weight_zero(self) -> None:
         """
         Remove user classes having a weight of zero.
         """
@@ -235,20 +235,20 @@ def _remove_user_classes_with_weight_zero(self):
             raise ValueError("There are no users with weight > 0.")
         self.user_classes[:] = filtered_user_classes
 
-    def assign_equal_weights(self):
+    def assign_equal_weights(self) -> None:
         """
         Update the user classes such that each user runs their specified tasks with equal
         probability.
         """
         for u in self.user_classes:
             u.weight = 1
-            user_tasks = []
+            user_tasks: List[Union[TaskSet, Callable]] = []
             tasks_frontier = u.tasks
             while len(tasks_frontier) != 0:
                 t = tasks_frontier.pop()
-                if hasattr(t, "tasks") and t.tasks:
+                if isinstance(t, TaskHolder):
                     tasks_frontier.extend(t.tasks)
-                elif isinstance(t, Callable):
+                elif callable(t):
                     if t not in user_tasks:
                         user_tasks.append(t)
                 else: