Annotate decide_worker's return type

dask · Dec 9, 2020 · f8b8c14 · f8b8c14
1 parent 3d493a6
commit f8b8c14
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -4653,7 +4653,7 @@ def transition_no_worker_waiting(self, key):
                 pdb.set_trace()
             raise
 
-    def decide_worker(self, ts: TaskState):
+    def decide_worker(self, ts: TaskState) -> WorkerState:
         """
         Decide on a worker for task *ts*.  Return a WorkerState.
         """
@@ -6082,7 +6082,9 @@ def adaptive_target(self, comm=None, target_duration=None):
             return len(self.workers) - len(to_close)
 
 
-def decide_worker(ts: TaskState, all_workers, valid_workers: set, objective):
+def decide_worker(
+    ts: TaskState, all_workers, valid_workers: set, objective
+) -> WorkerState:
     """
     Decide which worker should take task *ts*.