zict type annotations

.pre-commit-config.yaml

@@ -42,3 +42,6 @@ repos:
           - numpy
           - dask
           - tornado
+          # TODO replace before merging
+          #- zict
+          - git+https://github.com/crusaderky/zict@annotations

continuous_integration/environment-3.9.yaml

@@ -52,7 +52,9 @@ dependencies:
   - pip:
       - git+https://github.com/dask/dask
       - git+https://github.com/dask/s3fs
-      - git+https://github.com/dask/zict
+      # TODO revert before merging
+      # - git+https://github.com/dask/zict
+      - git+https://github.com/crusaderky/zict@annotations
       # FIXME https://github.com/dask/distributed/issues/5345
       # - git+https://github.com/intake/filesystem_spec
       - git+https://github.com/joblib/joblib

distributed/spill.py

@@ -5,7 +5,7 @@
 from collections.abc import Mapping
 from contextlib import contextmanager
 from functools import partial
-from typing import Any, Literal, NamedTuple
+from typing import Any, Literal, NamedTuple, cast
 
 import zict
 from packaging.version import parse as parse_version
@@ -32,6 +32,7 @@ def __sub__(self, other: SpilledSize) -> SpilledSize:  # type: ignore
         return SpilledSize(self.memory - other.memory, self.disk - other.disk)
 
 
+# zict.buffer[str, Any] requires zict >= 2.2.0
 class SpillBuffer(zict.Buffer):
     """MutableMapping that automatically spills out dask key/value pairs to disk when
     the total size of the stored data exceeds the target. If max_spill is provided the
@@ -171,7 +172,7 @@ def evict(self) -> int:
         try:
             with self.handle_errors(None):
                 _, _, weight = self.fast.evict()
-                return weight
+                return cast(int, weight)
         except HandledError:
             return -1
 
@@ -203,7 +204,7 @@ def spilled_total(self) -> SpilledSize:
         The two may differ substantially, e.g. if sizeof() is inaccurate or in case of
         compression.
         """
-        return self.slow.total_weight
+        return cast(Slow, self.slow).total_weight
 
 
 def _in_memory_weight(key: str, value: Any) -> int:
@@ -240,7 +241,8 @@ def __init__(self, spill_directory: str, max_weight: int | Literal[False] = Fals
 
     def __setitem__(self, key: str, value: Any) -> None:
         try:
-            pickled = self.dump(value)
+            # FIXME https://github.com/python/mypy/issues/708
+            pickled = self.dump(value)  # type: ignore
         except Exception as e:
             # zict.LRU ensures that the key remains in fast if we raise.
             # Wrap the exception so that it's recognizable by SpillBuffer,

distributed/tests/test_worker.py

@@ -73,7 +73,7 @@
 try:
     import zict
 except ImportError:
-    zict = None
+    zict = None  # type: ignore
 
 requires_zict = pytest.mark.skipif(not zict, reason="requires zict")
 requires_zict_210 = pytest.mark.skipif(