(refactor): remove read_dispatched_async

Author: ilan-gold

src/anndata/_io/h5ad.py

@@ -24,7 +24,7 @@
     _decode_structured_array,
     _from_fixed_length_strings,
 )
-from ..experimental import read_dispatched_async
+from ..experimental import read_dispatched
 from .specs import write_elem
 from .specs.methods import sync_async_to_async
 from .specs.registry import IOSpec, read_elem_async, write_spec
@@ -243,9 +243,7 @@ async def callback(func, elem_name: str, elem, iospec):
                         *(
                             # This is covering up backwards compat in the anndata initializer
                             # In most cases we should be able to call `func(elen[k])` instead
-                            sync_async_to_async(
-                                k, read_dispatched_async(elem[k], callback)
-                            )
+                            sync_async_to_async(k, read_dispatched(elem[k], callback))
                             for k in elem.keys()
                             if not k.startswith("raw.")
                         )
@@ -263,7 +261,7 @@ async def callback(func, elem_name: str, elem, iospec):
                 return await read_dataframe(elem)
             return await func(elem)
 
-        adata = asyncio.run(read_dispatched_async(f, callback=callback))
+        adata = asyncio.run(read_dispatched(f, callback=callback))
 
         # Backwards compat (should figure out which version)
         if "raw.X" in f:

src/anndata/_io/specs/methods.py

@@ -313,12 +313,6 @@ def write_null_zarr(f, k, _v, _writer, dataset_kwargs=MappingProxyType({})):
 async def read_mapping(
     elem: GroupStorageType, *, _reader: Reader
 ) -> dict[str, AxisStorable]:
-    print(
-        (
-            sync_async_to_async(k, _reader.read_elem_async(v))
-            for k, v in dict(elem).items()
-        )
-    )
     return dict(
         await asyncio.gather(
             *(

src/anndata/_io/zarr.py

@@ -14,7 +14,7 @@
 from .._settings import settings
 from .._warnings import OldFormatWarning
 from ..compat import _clean_uns, _from_fixed_length_strings, is_zarr_v2
-from ..experimental import read_dispatched_async, write_dispatched
+from ..experimental import read_dispatched, write_dispatched
 from .specs import read_elem_async
 from .specs.methods import sync_async_to_async
 from .utils import _read_legacy_raw, report_read_key_on_error
@@ -84,7 +84,7 @@ async def callback(func, elem_name: str, elem, iospec):
                     *(
                         # This is covering up backwards compat in the anndata initializer
                         # In most cases we should be able to call `func(elen[k])` instead
-                        sync_async_to_async(k, read_dispatched_async(elem[k], callback))
+                        sync_async_to_async(k, read_dispatched(elem[k], callback))
                         for k in elem.keys()
                         if not k.startswith("raw.")
                     )
@@ -102,7 +102,7 @@ async def callback(func, elem_name: str, elem, iospec):
             )
         return await func(elem)
 
-    adata = asyncio.run(read_dispatched_async(f, callback=callback))
+    adata = asyncio.run(read_dispatched(f, callback=callback))
 
     # Backwards compat (should figure out which version)
     if "raw.X" in f:

src/anndata/experimental/__init__.py

@@ -6,7 +6,7 @@
 from .._io.specs import IOSpec, read_elem_as_dask
 from .._types import Read, ReadCallback, StorageType, Write, WriteCallback
 from ..utils import module_get_attr_redirect
-from ._dispatch_io import read_dispatched, read_dispatched_async, write_dispatched
+from ._dispatch_io import read_dispatched, write_dispatched
 from .merge import concat_on_disk
 from .multi_files import AnnCollection
 from .pytorch import AnnLoader
@@ -43,7 +43,6 @@ def __getattr__(attr_name: str) -> Any:
     "AnnLoader",
     "read_elem_as_dask",
     "read_dispatched",
-    "read_dispatched_async",
     "write_dispatched",
     "IOSpec",
     "concat_on_disk",

src/anndata/experimental/_dispatch_io.py

@@ -10,40 +10,13 @@
     from anndata._types import (
         GroupStorageType,
         ReadAsyncCallback,
-        ReadCallback,
         StorageType,
         WriteCallback,
     )
     from anndata.typing import RWAble
 
 
-def read_dispatched(
-    elem: StorageType,
-    callback: ReadCallback,
-) -> RWAble:
-    """
-    Read elem, calling the callback at each sub-element.
-
-    Params
-    ------
-    elem
-        Storage container (e.g. `h5py.Group`, `zarr.Group`).
-        This must have anndata element specifications.
-    callback
-        Function to call at each anndata encoded element.
-
-    See Also
-    --------
-    :doc:`/tutorials/notebooks/{read,write}_dispatched`
-    """
-    from anndata._io.specs import _REGISTRY, Reader
-
-    reader = Reader(_REGISTRY, callback=callback)
-
-    return reader.read_elem(elem)
-
-
-async def read_dispatched_async(
+async def read_dispatched(
     elem: StorageType,
     callback: ReadAsyncCallback,
 ) -> RWAble:

src/anndata/experimental/merge.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import os
 import shutil
 from collections.abc import Mapping
@@ -26,6 +27,8 @@
 )
 from .._core.sparse_dataset import BaseCompressedSparseDataset, sparse_dataset
 from .._io.specs import read_elem, write_elem
+from .._io.specs.methods import sync_async_to_async
+from .._io.specs.registry import read_elem_async
 from ..compat import H5Array, H5Group, ZarrArray, ZarrGroup, _map_cat_to_str
 from . import read_dispatched
 
@@ -142,19 +145,26 @@ def read_as_backed(group: ZarrGroup | H5Group):
     BaseCompressedSparseDataset, Array or EAGER_TYPES are encountered.
     """
 
-    def callback(func, elem_name: str, elem, iospec):
+    async def callback(func, elem_name: str, elem, iospec):
         if iospec.encoding_type in SPARSE_MATRIX:
             return sparse_dataset(elem)
         elif iospec.encoding_type in EAGER_TYPES:
-            return read_elem(elem)
+            return await read_elem_async(elem)
         elif iospec.encoding_type == "array":
             return elem
         elif iospec.encoding_type == "dict":
-            return {k: read_as_backed(v) for k, v in dict(elem).items()}
+            return dict(
+                await asyncio.gather(
+                    *(
+                        sync_async_to_async(k, read_dispatched(v, callback=callback))
+                        for k, v in dict(elem).items()
+                    )
+                )
+            )
         else:
-            return func(elem)
+            return await func(elem)
 
-    return read_dispatched(group, callback=callback)
+    return asyncio.run(read_dispatched(group, callback=callback))
 
 
 def _df_index(df: ZarrGroup | H5Group) -> pd.Index:

src/anndata/io.py

@@ -14,7 +14,7 @@
     read_text,
     read_umi_tools,
 )
-from ._io.specs import read_elem, write_elem
+from ._io.specs import read_elem, read_elem_async, write_elem
 from ._io.write import write_csvs, write_loom
 
 if find_spec("zarr") or TYPE_CHECKING:
@@ -46,5 +46,6 @@ def write_zarr(*args, **kw):
     "write_zarr",
     "write_elem",
     "read_elem",
+    "read_elem_async",
     "sparse_dataset",
 ]

tests/test_backed_sparse.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from functools import partial
 from itertools import product
 from typing import TYPE_CHECKING, Literal, get_args
@@ -84,19 +85,12 @@ def read_zarr_backed(path):
             f = zarr.open(path, mode="r")
 
             # Read with handling for backwards compat
-            def callback(func, elem_name, elem, iospec):
-                if iospec.encoding_type == "anndata" or elem_name.endswith("/"):
-                    return AnnData(
-                        **{
-                            k: read_dispatched(v, callback)
-                            for k, v in dict(elem).items()
-                        }
-                    )
+            async def callback(func, elem_name, elem, iospec):
                 if iospec.encoding_type in {"csc_matrix", "csr_matrix"}:
                     return sparse_dataset(elem)
-                return func(elem)
+                return await func(elem)
 
-            adata = read_dispatched(f, callback=callback)
+            adata = asyncio.run(read_dispatched(f, callback=callback))
 
             return adata
 

tests/test_io_dispatched.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import re
 from typing import TYPE_CHECKING
 
@@ -18,11 +19,11 @@
 
 
 def test_read_dispatched_w_regex(tmp_path: Path):
-    def read_only_axis_dfs(func, elem_name: str, elem, iospec):
+    async def read_only_axis_dfs(func, elem_name: str, elem, iospec):
         if iospec.encoding_type == "anndata":
-            return func(elem)
+            return await func(elem)
         elif re.match(r"^/((obs)|(var))?(/.*)?$", elem_name):
-            return func(elem)
+            return await func(elem)
         else:
             return None
 
@@ -35,27 +36,27 @@ def read_only_axis_dfs(func, elem_name: str, elem, iospec):
         z = zarr.open(z.store)
 
     expected = ad.AnnData(obs=adata.obs, var=adata.var)
-    actual = read_dispatched(z, read_only_axis_dfs)
+    actual = asyncio.run(read_dispatched(z, read_only_axis_dfs))
 
     assert_equal(expected, actual)
 
 
 def test_read_dispatched_dask(tmp_path: Path):
     import dask.array as da
 
-    def read_as_dask_array(func, elem_name: str, elem, iospec):
+    async def read_as_dask_array(func, elem_name: str, elem, iospec):
         if iospec.encoding_type in {
             "dataframe",
             "csr_matrix",
             "csc_matrix",
             "awkward-array",
         }:
             # Preventing recursing inside of these types
-            return func(elem)
+            return await func(elem)
         elif iospec.encoding_type == "array":
             return da.from_zarr(elem)
         else:
-            return func(elem)
+            return await func(elem)
 
     adata = gen_adata((1000, 100))
     z = open_write_group(tmp_path)
@@ -64,7 +65,7 @@ def read_as_dask_array(func, elem_name: str, elem, iospec):
     if not is_zarr_v2() and isinstance(z, ZarrGroup):
         z = zarr.open(z.store)
 
-    dask_adata = read_dispatched(z, read_as_dask_array)
+    dask_adata = asyncio.run(read_dispatched(z, read_as_dask_array))
 
     assert isinstance(dask_adata.layers["array"], da.Array)
     assert isinstance(dask_adata.obsm["array"], da.Array)
@@ -84,7 +85,11 @@ def test_read_dispatched_null_case(tmp_path: Path):
     if not is_zarr_v2() and isinstance(z, ZarrGroup):
         z = zarr.open(z.store)
     expected = ad.io.read_elem(z)
-    actual = read_dispatched(z, lambda _, __, x, **___: ad.io.read_elem(x))
+
+    async def callback(_, __, x, **___):
+        return await ad.io.read_elem_async(x)
+
+    actual = asyncio.run(read_dispatched(z, callback))
 
     assert_equal(expected, actual)
 
@@ -186,23 +191,23 @@ def zarr_writer(func, store, k, elem, dataset_kwargs, iospec):
         )
         func(store, k, elem, dataset_kwargs=dataset_kwargs)
 
-    def h5ad_reader(func, elem_name: str, elem, iospec):
+    async def h5ad_reader(func, elem_name: str, elem, iospec):
         h5ad_read_keys.append(elem_name if is_zarr_v2() else elem_name.strip("/"))
         return func(elem)
 
-    def zarr_reader(func, elem_name: str, elem, iospec):
+    async def zarr_reader(func, elem_name: str, elem, iospec):
         zarr_read_keys.append(elem_name if is_zarr_v2() else elem_name.strip("/"))
         return func(elem)
 
     adata = gen_adata((50, 100))
 
     with h5py.File(h5ad_path, "w") as f:
         write_dispatched(f, "/", adata, callback=h5ad_writer)
-        _ = read_dispatched(f, h5ad_reader)
+        _ = asyncio.run(read_dispatched(f, h5ad_reader))
 
     f = open_write_group(zarr_path)
     write_dispatched(f, "/", adata, callback=zarr_writer)
-    _ = read_dispatched(f, zarr_reader)
+    _ = asyncio.run(read_dispatched(f, zarr_reader))
 
     assert sorted(h5ad_read_keys) == sorted(zarr_read_keys)
     assert sorted(h5ad_write_keys) == sorted(zarr_write_keys)