From b94310fb276304595e6f4f2c644233f8092395cc Mon Sep 17 00:00:00 2001
From: John Kirkham <jakirkham@gmail.com>
Date: Tue, 8 Dec 2020 17:58:29 -0800
Subject: [PATCH] Annotate `decide_worker`'s return type

---
 distributed/scheduler.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index a0684ed3b7..0d7d48f827 100644
--- 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*.