Backport PR #2433 on branch 0.12.x (fix: correct IO types) (#2434)

Author: flying-sheep

.vscode/settings.json

@@ -24,4 +24,7 @@
         //"-nauto",
     ],
     "python.terminal.activateEnvironment": true,
+    "python.analysis.include": ["src/**/*", "ci/scripts/**/*", "tests/**/*"],
+    "python-envs.defaultEnvManager": "pypa.hatch:hatch",
+    "python-envs.defaultPackageManager": "pypa.hatch:hatch",
 }

src/anndata/_io/specs/lazy_methods.py

@@ -202,10 +202,7 @@ def resolve_chunks(
 # In the long run, it might be good to figure out what exactly is going on here but for now, this will do.
 @_LAZY_REGISTRY.register_read(H5Array, IOSpec("string-array", "0.2.0"))
 def read_h5_string_array(
-    elem: H5Array,
-    *,
-    _reader: LazyReader,
-    chunks: tuple[int] | None = None,
+    elem: H5Array, *, _reader: LazyReader, chunks: tuple[int, ...] | None = None
 ) -> DaskArray:
     import dask.array as da
 
@@ -249,7 +246,7 @@ def read_zarr_array(
 def _gen_xarray_dict_iterator_from_elems(
     elem_dict: dict[str, LazyDataStructures],
     dim_name: str,
-    index: np.NDArray,
+    index: np.typing.NDArray,
 ) -> Generator[tuple[str, XVariable], None, None]:
     from anndata.experimental.backed._lazy_arrays import CategoricalArray, MaskedArray
 
@@ -291,7 +288,7 @@ def read_dataframe(
     *,
     _reader: LazyReader,
     use_range_index: bool = False,
-    chunks: tuple[int] | None = None,
+    chunks: tuple[int, ...] | None = None,
 ) -> Dataset2D:
     # going through dask for reading into memory the index doesn't make sense, hence the ternary.
     elem_dict = {
@@ -337,9 +334,12 @@ def read_categorical(
     elem: H5Group | ZarrGroup,
     *,
     _reader: LazyReader,
+    chunks: tuple[int, ...] | None = None,
 ) -> CategoricalArray:
     from anndata.experimental.backed._lazy_arrays import CategoricalArray
 
+    del chunks  # ignored when reading groups
+
     base_path_or_zarr_group = (
         Path(filename(elem)) if isinstance(elem, H5Group) else elem
     )
@@ -361,9 +361,12 @@ def read_nullable(
         "nullable-integer", "nullable-boolean", "nullable-string-array"
     ],
     _reader: LazyReader,
+    chunks: tuple[int, ...] | None = None,
 ) -> MaskedArray:
     from anndata.experimental.backed._lazy_arrays import MaskedArray
 
+    del chunks  # ignored when reading groups
+
     base_path_or_zarr_group = (
         Path(filename(elem)) if isinstance(elem, H5Group) else elem
     )

src/anndata/_io/specs/methods.py

@@ -7,7 +7,7 @@
 from importlib.metadata import version
 from itertools import product
 from types import MappingProxyType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 from warnings import warn
 
 import h5py
@@ -46,7 +46,7 @@
 from .registry import _REGISTRY, IOSpec, read_elem, read_elem_partial
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterator
+    from collections.abc import Iterator
     from os import PathLike
     from typing import Any, Literal
 
@@ -443,7 +443,7 @@ def write_basic(
     dataset_kwargs: Mapping[str, Any] = MappingProxyType({}),
 ):
     """Write methods which underlying library handles natively."""
-    dataset_kwargs = dataset_kwargs.copy()
+    dataset_kwargs = dict(dataset_kwargs)
     dtype = dataset_kwargs.pop("dtype", elem.dtype)
     if isinstance(f, H5Group) or is_zarr_v2():
         f.create_dataset(k, data=elem, shape=elem.shape, dtype=dtype, **dataset_kwargs)
@@ -525,7 +525,7 @@ def write_basic_dask_dask_dense(
 ):
     import dask.array as da
 
-    dataset_kwargs = dataset_kwargs.copy()
+    dataset_kwargs = dict(dataset_kwargs)
     is_h5 = isinstance(f, H5Group)
     if not is_h5:
         dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
@@ -629,7 +629,7 @@ def write_vlen_string_array_zarr(
         from numcodecs import VLenUTF8
         from zarr.core.dtype import VariableLengthUTF8
 
-        dataset_kwargs = dataset_kwargs.copy()
+        dataset_kwargs = dict(dataset_kwargs)
         dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
         dtype = VariableLengthUTF8()
         filters, fill_value = None, None
@@ -702,7 +702,7 @@ def write_recarray_zarr(
     if is_zarr_v2():
         f.create_dataset(k, data=elem, shape=elem.shape, **dataset_kwargs)
     else:
-        dataset_kwargs = dataset_kwargs.copy()
+        dataset_kwargs = dict(dataset_kwargs)
         dataset_kwargs = zarr_v3_compressor_compat(dataset_kwargs)
         # https://github.com/zarr-developers/zarr-python/issues/3546
         # if "shards" not in dataset_kwargs and ad.settings.auto_shard_zarr_v3:
@@ -1210,14 +1210,14 @@ def write_nullable(
     )(write_nullable)
 
 
+class _BaseMaskedArray(Protocol):
+    def __call__(
+        self, values: NDArray[np.number], /, *, mask: NDArray[np.bool_]
+    ) -> pd.api.extensions.ExtensionArray: ...
+
+
 def _read_nullable(
-    elem: GroupStorageType,
-    *,
-    _reader: Reader,
-    # BaseMaskedArray
-    array_type: Callable[
-        [NDArray[np.number], NDArray[np.bool_]], pd.api.extensions.ExtensionArray
-    ],
+    elem: GroupStorageType, *, _reader: Reader, array_type: _BaseMaskedArray
 ) -> pd.api.extensions.ExtensionArray:
     return array_type(
         _reader.read_elem(elem["values"]),
@@ -1378,7 +1378,7 @@ def write_string(
     _writer: Writer,
     dataset_kwargs: Mapping[str, Any],
 ):
-    dataset_kwargs = dataset_kwargs.copy()
+    dataset_kwargs = dict(dataset_kwargs)
     dataset_kwargs.pop("compression", None)
     dataset_kwargs.pop("compression_opts", None)
     f.create_dataset(

src/anndata/_io/specs/registry.py

@@ -30,6 +30,7 @@
 
     from ..._core.xarray import Dataset2D
 
+    S = TypeVar("S")
     T = TypeVar("T")
     W = TypeVar("W", bound=_WriteInternal)
     LazyDataStructures = DaskArray | Dataset2D | CategoricalArray | MaskedArray
@@ -103,7 +104,7 @@ def register_write(
         src_type: type | tuple[type, str],
         spec: IOSpec | Mapping[str, str],
         modifiers: Iterable[str] = frozenset(),
-    ) -> Callable[[_WriteInternal[T]], _WriteInternal[T]]:
+    ) -> Callable[[_WriteInternal[S, T]], _WriteInternal[S, T]]:
         spec = proc_spec(spec)
         modifiers = frozenset(modifiers)
 
@@ -119,7 +120,7 @@ def register_write(
         else:
             self.write_specs[src_type] = spec
 
-        def _register(func):
+        def _register(func: _WriteInternal[S, T]) -> _WriteInternal[S, T]:
             self.write[(dest_type, src_type, modifiers)] = write_spec(spec)(func)
             return func
 
@@ -351,7 +352,7 @@ def write_elem(
 
         # we allow stores to have a prefix like /uns which are then written to with keys like /uns/foo
         is_zarr_group = isinstance(store, ZarrGroup)
-        if "/" in k.split(store.name)[-1][1:]:
+        if "/" in k.rsplit(store.name, maxsplit=1)[-1][1:]:
             if is_zarr_group or settings.disallow_forward_slash_in_h5ad:
                 msg = f"Forward slashes are not allowed in keys in {type(store)}"
                 raise ValueError(msg)

src/anndata/_io/utils.py

@@ -312,7 +312,7 @@ def _read_legacy_raw(
     return raw
 
 
-def zero_dim_array_as_scalar(func: _WriteInternal):
+def zero_dim_array_as_scalar(func: _WriteInternal) -> _WriteInternal:
     """\
     A decorator for write_elem implementations of arrays where zero-dimensional arrays need special handling.
     """

src/anndata/_types.py

@@ -25,9 +25,16 @@
 
 
 __all__ = [
+    "AnnDataElem",
     "ArrayStorageType",
+    "Dataset2DIlocIndexer",
     "GroupStorageType",
+    "Read",
+    "ReadCallback",
+    "ReadLazy",
     "StorageType",
+    "Write",
+    "WriteCallback",
     "_ReadInternal",
     "_ReadLazyInternal",
     "_WriteInternal",
@@ -47,25 +54,26 @@
 
 
 class Dataset2DIlocIndexer(Protocol):
-    def __getitem__(self, idx: Any) -> Dataset2D: ...
+    def __getitem__(self, idx: Any, /) -> Dataset2D: ...
 
 
 class _ReadInternal(Protocol[S_contra, RWAble_co]):
-    def __call__(self, elem: S_contra, *, _reader: Reader) -> RWAble_co: ...
+    def __call__(self, elem: S_contra, /, *, _reader: Reader) -> RWAble_co: ...
 
 
 class _ReadLazyInternal(Protocol[S_contra]):
     def __call__(
         self,
         elem: S_contra,
+        /,
         *,
         _reader: LazyReader,
         chunks: tuple[int, ...] | None = None,
     ) -> LazyDataStructures: ...
 
 
 class Read(Protocol[S_contra, RWAble_co]):
-    def __call__(self, elem: S_contra) -> RWAble_co:
+    def __call__(self, elem: S_contra, /) -> RWAble_co:
         """Low-level reading function for an element.
 
         Parameters
@@ -81,7 +89,7 @@ def __call__(self, elem: S_contra) -> RWAble_co:
 
 class ReadLazy(Protocol[S_contra]):
     def __call__(
-        self, elem: S_contra, *, chunks: tuple[int, ...] | None = None
+        self, elem: S_contra, /, *, chunks: tuple[int, ...] | None = None
     ) -> LazyDataStructures:
         """Low-level reading function for a lazy element.
 
@@ -98,12 +106,13 @@ def __call__(
         ...
 
 
-class _WriteInternal(Protocol[RWAble_contra]):
+class _WriteInternal(Protocol[S_contra, RWAble_contra]):
     def __call__(
         self,
-        f: StorageType,
+        f: S_contra,
         k: str,
         v: RWAble_contra,
+        /,
         *,
         _writer: Writer,
         dataset_kwargs: Mapping[str, Any],
@@ -116,6 +125,7 @@ def __call__(
         f: StorageType,
         k: str,
         v: RWAble_contra,
+        /,
         *,
         dataset_kwargs: Mapping[str, Any],
     ) -> None: