Review scheduler annotations, part 2 (#6253)

dask · May 11, 2022 · a4999a9 · a4999a9
1 parent e390609
commit a4999a9
Show file tree

Hide file tree

Showing 4 changed files with 442 additions and 561 deletions.
diff --git a/distributed/batched.py b/distributed/batched.py
@@ -128,7 +128,7 @@ def _background_send(self):
         self.stopped.set()
         self.abort()
 
-    def send(self, *msgs: dict) -> None:
+    def send(self, *msgs) -> None:
         """Schedule a message for sending to the other side
 
         This completes quickly and synchronously

diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
@@ -72,7 +72,11 @@ async def close(self) -> None:
         """
 
     def update_graph(
-        self, scheduler: Scheduler, keys: set[str], restrictions: dict, **kwargs
+        self,
+        scheduler: Scheduler,
+        keys: set[str],
+        restrictions: dict[str, float],
+        **kwargs,
     ) -> None:
         """Run when a new graph / tasks enter the scheduler"""