Add generic type support to Future and Client methods (#9123)

Author: moi90

distributed/client.py

@@ -31,9 +31,11 @@
     Any,
     Callable,
     ClassVar,
+    Generic,
     Literal,
     NamedTuple,
     TypedDict,
+    TypeVar,
     cast,
 )
 
@@ -138,6 +140,8 @@
 
 TOPIC_PREFIX_FORWARDED_LOG_RECORD = "forwarded-log-record"
 
+_T = TypeVar("_T")
+
 
 class FutureCancelledError(CancelledError):
     key: str
@@ -241,7 +245,7 @@ def _del_global_client(c: Client) -> None:
             pass
 
 
-class Future(TaskRef):
+class Future(TaskRef, Generic[_T]):
     """A remotely running computation
 
     A Future is a local proxy to a result running on a remote worker.  A user
@@ -371,7 +375,7 @@ def done(self):
         """
         return self._state.done()
 
-    def result(self, timeout=None):
+    def result(self, timeout=None) -> _T:
         """Wait until computation completes, gather result to local process.
 
         Parameters
@@ -2033,7 +2037,7 @@ def get_executor(self, **kwargs):
 
     def submit(
         self,
-        func,
+        func: Callable[..., _T],
         *args,
         key=None,
         workers=None,
@@ -2046,7 +2050,7 @@ def submit(
         actors=False,
         pure=True,
         **kwargs,
-    ):
+    ) -> Future[_T]:
         """Submit a function application to the scheduler
 
         Parameters
@@ -2154,7 +2158,7 @@ def submit(
                     key,
                     func,
                     *(parse_input(a) for a in args),
-                    **{k: parse_input(v) for k, v in kwargs.items()},
+                    **{k: parse_input(v) for k, v in kwargs.items()},  # type: ignore
                 )
             },
             # We'd like to avoid hashing/tokenizing all of the above.
@@ -2181,7 +2185,7 @@ def submit(
 
     def map(
         self,
-        func: Callable,
+        func: Callable[..., _T],
         *iterables: Collection,
         key: str | list | None = None,
         workers: str | Iterable[str] | None = None,
@@ -2195,7 +2199,7 @@ def map(
         pure: bool = True,
         batch_size=None,
         **kwargs,
-    ):
+    ) -> list[Future[_T]]:
         """Map a function on a sequence of arguments
 
         Arguments can be normal objects or Futures