[v3] Fix some untyped defs

pyproject.toml

@@ -161,15 +161,11 @@ check_untyped_defs = true
 
 [[tool.mypy.overrides]]
 module = [
-    "zarr._storage.store",
     "zarr._storage.v3_storage_transformers",
     "zarr.v3.group",
     "zarr.core",
     "zarr.hierarchy",
-    "zarr.indexing",
     "zarr.storage",
-    "zarr.sync",
-    "zarr.util",
     "tests.*",
 ]
 check_untyped_defs = false

src/zarr/_storage/store.py

@@ -340,20 +340,20 @@ def set_partial_values(self, key_start_values):
             old_value = self.get(key)
             values[key] = None if old_value is None else bytearray(old_value)
         for key, start, value in key_start_values:
-            if values[key] is None:
+            if (old_value := values[key]) is None:
                 assert start == 0
                 values[key] = value
             else:
-                if start > len(values[key]):  # pragma: no cover
+                if start > len(old_value):
                     raise ValueError(
                         f"Cannot set value at start {start}, "
                         + f"since it is beyond the data at key {key}, "
-                        + f"having length {len(values[key])}."
+                        + f"having length {len(old_value)}."
                     )
                 if start < 0:
-                    values[key][start:] = value
+                    old_value[start:] = value
                 else:
-                    values[key][start : start + len(value)] = value
+                    old_value[start : start + len(value)] = value
         for key, value in values.items():
             self[key] = value
 
@@ -474,7 +474,7 @@ def list(self):
         return list(self.keys())
 
     def list_dir(self, prefix):
-        return StoreV3.list_dir(self, prefix)
+        return self.inner_store.list_dir(prefix)
 
     def is_readable(self):
         return self.inner_store.is_readable()

src/zarr/indexing.py

@@ -4,6 +4,7 @@
 import numbers
 
 import numpy as np
+import numpy.typing as npt
 
 
 from zarr.errors import (
@@ -330,6 +331,7 @@ def __init__(self, selection, array):
         # setup per-dimension indexers
         dim_indexers = []
         for dim_sel, dim_len, dim_chunk_len in zip(selection, array._shape, array._chunks):
+            dim_indexer: IntDimIndexer | SliceDimIndexer
             if is_integer(dim_sel):
                 dim_indexer = IntDimIndexer(dim_sel, dim_len, dim_chunk_len)
 
@@ -521,9 +523,12 @@ def __iter__(self):
             else:
                 start = self.chunk_nitems_cumsum[dim_chunk_ix - 1]
             stop = self.chunk_nitems_cumsum[dim_chunk_ix]
+
+            dim_out_sel: slice | np.ndarray
             if self.order == Order.INCREASING:
                 dim_out_sel = slice(start, stop)
             else:
+                assert self.dim_out_sel is not None
                 dim_out_sel = self.dim_out_sel[start:stop]
 
             # find region in chunk
@@ -577,11 +582,10 @@ def oindex_set(a, selection, value):
     selection = ix_(selection, a.shape)
     if not np.isscalar(value) and drop_axes:
         value = np.asanyarray(value)
-        value_selection = [slice(None)] * len(a.shape)
+        value_selection: list[slice | None] = [slice(None)] * len(a.shape)
         for i in drop_axes:
             value_selection[i] = np.newaxis
-        value_selection = tuple(value_selection)
-        value = value[value_selection]
+        value = value[tuple(value_selection)]
     a[selection] = value
 
 
@@ -597,6 +601,7 @@ def __init__(self, selection, array):
         # setup per-dimension indexers
         dim_indexers = []
         for dim_sel, dim_len, dim_chunk_len in zip(selection, array._shape, array._chunks):
+            dim_indexer: IntDimIndexer | SliceDimIndexer | IntArrayDimIndexer | BoolArrayDimIndexer
             if is_integer(dim_sel):
                 dim_indexer = IntDimIndexer(dim_sel, dim_len, dim_chunk_len)
 
@@ -622,6 +627,8 @@ def __init__(self, selection, array):
         self.dim_indexers = dim_indexers
         self.shape = tuple(s.nitems for s in self.dim_indexers if not isinstance(s, IntDimIndexer))
         self.is_advanced = not is_basic_selection(selection)
+
+        self.drop_axes: tuple[int, ...] | None
         if self.is_advanced:
             self.drop_axes = tuple(
                 i
@@ -797,6 +804,7 @@ def __init__(self, selection, array):
             boundscheck_indices(dim_sel, dim_len)
 
         # compute chunk index for each point in the selection
+        chunks_multi_index: npt.ArrayLike
         chunks_multi_index = tuple(
             dim_sel // dim_chunk_len for (dim_sel, dim_chunk_len) in zip(selection, array._chunks)
         )
@@ -848,6 +856,8 @@ def __iter__(self):
             else:
                 start = self.chunk_nitems_cumsum[chunk_rix - 1]
             stop = self.chunk_nitems_cumsum[chunk_rix]
+
+            out_selection: slice | np.ndarray
             if self.sel_sort is None:
                 out_selection = slice(start, stop)
             else:
@@ -952,6 +962,7 @@ def check_no_multi_fields(fields):
 
 
 def pop_fields(selection):
+    fields: str | None | list[str]
     if isinstance(selection, str):
         # single field selection
         fields = selection

src/zarr/sync.py

@@ -21,7 +21,8 @@ def __getstate__(self):
 
     def __setstate__(self, *args):
         # reinitialize from scratch
-        self.__init__()
+        self.mutex = Lock()
+        self.locks = defaultdict(Lock)
 
 
 class ProcessSynchronizer:

src/zarr/util.py

@@ -623,6 +623,9 @@ def prepare_chunk(self):
 
     def read_part(self, start, nitems):
         assert self.buff is not None
+        assert self.start_points is not None
+        assert self.n_per_block is not None
+
         if self.nblocks == 1:
             return
         start_block = int(start / self.n_per_block)