Repository: locustio/locust

Issue:
Add type hints
<!-- 
If you have a general question about how to use Locust, please check Stack Overflow first https://stackoverflow.com/questions/tagged/locust

You can also ask new questions on SO, https://stackoverflow.com/questions/ask just remember to tag your question with "locust".
-->

### Is your feature request related to a problem? Please describe.
<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] -->

Type hints are just a general improvement :)

### Describe the solution you'd like
<!-- A clear and concise description of what you want to happen -->

Provide in-line type hints and mark the package as providing types by including a `py.typed` file in the distribution.

### Additional context

You could also provide type stubs as pyright can automatically generate them for you, removing some of the work required for this:

```
pyright --createstub locust
```


Discussion:
Oh, are they not visible when doing a package install? (as opposed to a source install)

Or do you mean a type hint for a specific thing?

Either way, PRs welcome!
Ah sorry I should've clarified my situation. I use pyright which complains if a library doesn't provide in-line type hints or type stubs. Pyright only considers a library to provide in-line type hints if it has a `py.typed` file in the root directory of the library because of [PEP 561](https://www.python.org/dev/peps/pep-0561/#packaging-type-information).

I also ran into some parts of the public API that didn't have type hints (like `@task`) which pyright also doesn't like because I'm using strict mode.

I fixed these errors by running `pyright --createstub locust` and then manually adding some missing type hints to the parts of the interface that I'm using.

I'll have a look through the code base and if I have time I'll create a PR to add some missing type hints and provide the `py.typed` file :)

Resolving patch:
diff --git a/locust/user/task.py b/locust/user/task.py
--- a/locust/user/task.py
+++ b/locust/user/task.py
@@ -1,9 +1,8 @@
 import logging
 import random
-import sys
 import traceback
 from time import time
-from typing import Any, Callable, List, Union
+from typing import TYPE_CHECKING, Callable, Any, List, Union, TypeVar, overload
 from typing_extensions import final
 
 import gevent
@@ -11,13 +10,27 @@
 
 from locust.exception import InterruptTaskSet, RescheduleTask, RescheduleTaskImmediately, StopUser, MissingWaitTimeError
 
+if TYPE_CHECKING:
+    from locust import User
+
 
 logger = logging.getLogger(__name__)
+TaskT = TypeVar("TaskT", bound=Callable[[Any], None])
 
 LOCUST_STATE_RUNNING, LOCUST_STATE_WAITING, LOCUST_STATE_STOPPING = ["running", "waiting", "stopping"]
 
 
-def task(weight=1):
+@overload
+def task(weight: TaskT) -> TaskT:
+    ...
+
+
+@overload
+def task(weight: int) -> Callable[[TaskT], TaskT]:
+    ...
+
+
+def task(weight: Union[TaskT, int] = 1) -> Union[TaskT, Callable[[TaskT], TaskT]]:
     """
     Used as a convenience decorator to be able to declare tasks for a User or a TaskSet
     inline in the class. Example::
@@ -59,7 +72,7 @@ def my_task()
         return decorator_func
 
 
-def tag(*tags):
+def tag(*tags: str) -> Callable[[TaskT], TaskT]:
     """
     Decorator for tagging tasks and TaskSets with the given tag name. You can
     then limit the test to only execute tasks that are tagged with any of the
@@ -233,7 +246,7 @@ class ForumPage(TaskSet):
     _user = None
     _parent = None
 
-    def __init__(self, parent):
+    def __init__(self, parent: "User") -> None:
         self._task_queue = []
         self._time_start = time()
 
@@ -253,7 +266,7 @@ def __init__(self, parent):
             self.wait_function = self.user.wait_function
 
     @property
-    def user(self):
+    def user(self) -> "User":
         """:py:class:`User <locust.User>` instance that this TaskSet was created by"""
         return self._user
 


Test patch: