Enable mypy in CI 2/2

.pre-commit-config.yaml

@@ -4,6 +4,14 @@ repos:
       hooks:
       - id: isort
         language_version: python3
+  - repo: https://github.com/asottile/pyupgrade
+    # Do not upgrade: there's a bug in Cython that causes sum(... for ...) to fail;
+    # it needs sum([... for ...])
+    rev: v2.13.0
+    hooks:
+      - id: pyupgrade
+        args:
+          - --py37-plus
   -   repo: https://github.com/psf/black
       rev: 21.9b0
       hooks:
@@ -17,11 +25,22 @@ repos:
       hooks:
       - id: flake8
         language_version: python3
-  - repo: https://github.com/asottile/pyupgrade
-    # Do not upgrade: there's a bug in Cython that causes sum(... for ...) to fail;
-    # it needs sum([... for ...])
-    rev: v2.13.0
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.910
     hooks:
-      - id: pyupgrade
-        args:
-          - --py37-plus
+      - id: mypy
+        additional_dependencies:
+          # Type stubs
+          - types-docutils
+          - types-requests
+          - types-paramiko
+          - types-pkg_resources
+          - types-PyYAML
+          - types-setuptools
+          - types-psutil
+          # Libraries exclusively imported under `if TYPE_CHECKING:`
+          - typing_extensions  # To be reviewed after dropping Python 3.7
+          # Typed libraries
+          - numpy
+          - dask
+          - tornado

distributed/active_memory_manager.py

@@ -127,12 +127,12 @@ def run_once(self, comm=None) -> None:
                 # populate self.pending
                 self._run_policies()
 
-                drop_by_worker: defaultdict[WorkerState, set[TaskState]] = defaultdict(
-                    set
-                )
-                repl_by_worker: defaultdict[
-                    WorkerState, dict[TaskState, set[WorkerState]]
-                ] = defaultdict(dict)
+                drop_by_worker: (
+                    defaultdict[WorkerState, set[TaskState]]
+                ) = defaultdict(set)
+                repl_by_worker: (
+                    defaultdict[WorkerState, dict[TaskState, set[str]]]
+                ) = defaultdict(dict)
 
                 for ts, (pending_repl, pending_drop) in self.pending.items():
                     if not ts.who_has: