Simplify Generator annotations - PEP 696

examples/rest.py

@@ -77,7 +77,7 @@ class RestUserThatLooksAtErrors(FastHttpUser):
     abstract = True
 
     @contextmanager
-    def rest(self, method, url, **kwargs) -> Generator[RestResponseContextManager, None, None]:
+    def rest(self, method, url, **kwargs) -> Generator[RestResponseContextManager]:
         extra_headers = {"my_header": "my_value"}
         with super().rest(method, url, headers=extra_headers, **kwargs) as resp:
             if resp.js and "error" in resp.js and resp.js["error"] is not None:

locust/clients.py

@@ -122,7 +122,7 @@ def _build_url(self, path) -> str:
             return f"{self.base_url}{path}"
 
     @contextmanager
-    def rename_request(self, name: str) -> Generator[None, None, None]:
+    def rename_request(self, name: str) -> Generator[None]:
         """Group requests using the "with" keyword"""
 
         self.request_name = name

locust/contrib/fasthttp.py

@@ -407,9 +407,7 @@ def __init__(self, environment) -> None:
         """
 
     @contextmanager
-    def rest(
-        self, method, url, headers: dict | None = None, **kwargs
-    ) -> Generator[RestResponseContextManager, None, None]:
+    def rest(self, method, url, headers: dict | None = None, **kwargs) -> Generator[RestResponseContextManager]:
         """
         A wrapper for self.client.request that:
 
@@ -457,7 +455,7 @@ def rest(
                     resp.failure(f"{e.__class__.__name__}: {e} at {', '.join(error_lines)}. Response was {short_resp}")
 
     @contextmanager
-    def rest_(self, method, url, name=None, **kwargs) -> Generator[RestResponseContextManager, None, None]:
+    def rest_(self, method, url, name=None, **kwargs) -> Generator[RestResponseContextManager]:
         """
         Some REST api:s use a timestamp as part of their query string (mainly to break through caches).
         This is a convenience method for that, appending a _=<timestamp> parameter automatically

locust/dispatch.py

@@ -23,7 +23,7 @@
     T = TypeVar("T")
 
 
-def _kl_generator(users: Iterable[tuple[T, float]]) -> Iterator[T | None]:
+def _kl_generator(users: Iterable[tuple[T, float]]) -> Generator[T | None]:
     """Generator based on Kullback-Leibler divergence
 
     For example, given users A, B with weights 5 and 1 respectively,
@@ -97,7 +97,7 @@ def __init__(self, worker_nodes: list[WorkerNode], user_classes: list[type[User]
 
         self._current_user_count = self.get_current_user_count()
 
-        self._dispatcher_generator: Generator[dict[str, dict[str, int]], None, None] = None  # type: ignore
+        self._dispatcher_generator: Generator[dict[str, dict[str, int]]] = None  # type: ignore
         # a generator is assigned (in new_dispatch()) to _dispatcher_generator before it's used
 
         self._user_generator = self._user_gen()
@@ -149,7 +149,7 @@ def _sort_workers(self):
         # 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]]]:
         self._dispatch_in_progress = True
 
         if self._rebalance:
@@ -268,7 +268,7 @@ def _prepare_rebalance(self) -> None:
         self._rebalance = True
 
     @contextlib.contextmanager
-    def _wait_between_dispatch_iteration_context(self) -> Generator[None, None, None]:
+    def _wait_between_dispatch_iteration_context(self) -> Generator[None]:
         t0_rel = time.perf_counter()
 
         # We don't use `try: ... finally: ...` because we don't want to sleep

locust/event.py

@@ -56,7 +56,7 @@ def fire(self, *, reverse=False, **kwargs):
     @contextmanager
     def measure(
         self, request_type: str, name: str, response_length: int = 0, context=None
-    ) -> Generator[dict[str, Any], None, None]:
+    ) -> Generator[dict[str, Any]]:
         """Convenience method for firing the event with automatically calculated response time and automatically marking the request as failed if an exception is raised (this is really only useful for the *request* event)
 
         Example usage (in a task):