feat: AnnData.unwriteable based on AnnData._reduce + iter_outer + refactorings of other relevant functions

docs/api.md

@@ -264,6 +264,13 @@ Types used by the former:
    abc.CSCDataset
 ```
 
+```{eval-rst}
+.. autosummary::
+   :toctree: generated/
+
+   types.FoldFunc
+```
+
 <!-- these are types, not classes, so don’t use the above template -->
 
 ```{eval-rst}

docs/release-notes/2327.feat.md

@@ -0,0 +1 @@
+New {meth}`AnnData.fold` for crawling the "elems" and accumulating a value over these, and then {meth}`AnnData.can_write` built on top {user}`ilan-gold`

src/anndata/_core/anndata.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from collections.abc import Mapping, MutableMapping, Sequence
 from copy import copy, deepcopy
 from functools import partial, singledispatchmethod
@@ -22,6 +22,7 @@
 from scipy.sparse import issparse
 
 from anndata._warnings import ImplicitModificationWarning
+from anndata.acc import A, AdRef, GraphAcc, LayerAcc, MultiAcc
 
 from .. import utils
 from .._settings import settings
@@ -63,7 +64,10 @@
 
     from zarr.storage import StoreLike
 
-    from ..acc import AdRef, Array, MapAcc, RefAcc
+    from anndata.types import FoldFunc
+    from anndata.typing import RWAble
+
+    from ..acc import Array, MapAcc, RefAcc
     from ..compat import XDataset
     from ..typing import Index, Index1D, _Index1DNorm, _XDataType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
@@ -513,36 +517,39 @@ def _init_as_actual(  # noqa: PLR0912, PLR0913, PLR0915
     def __sizeof__(
         self, *, show_stratified: bool = False, with_disk: bool = False
     ) -> int:
-        def get_size(X) -> int:
+        def get_size[R: dict[RefAcc | None, int]](
+            X: RWAble,
+            *,
+            accumulate: R,
+            ref_acc: RefAcc | AdRef | None,
+        ) -> R:
             def cs_to_bytes(X) -> int:
                 return int(X.data.nbytes + X.indptr.nbytes + X.indices.nbytes)
 
-            if isinstance(X, h5py.Dataset) and with_disk:
-                return int(np.array(X.shape).prod() * X.dtype.itemsize)
-            elif isinstance(X, BaseCompressedSparseDataset) and with_disk:
-                return cs_to_bytes(X._to_backed())
-            elif issparse(X):
-                return cs_to_bytes(X)
-            else:
-                return X.__sizeof__()
-
-        sizes = {}
-        attrs = ["X", "_obs", "_var"]
-        attrs_multi = ["_uns", "_obsm", "_varm", "varp", "_obsp", "_layers"]
-        for attr in attrs + attrs_multi:
-            if attr in attrs_multi:
-                keys = getattr(self, attr).keys()
-                s = sum(get_size(getattr(self, attr)[k]) for k in keys)
-            else:
-                s = get_size(getattr(self, attr))
-            if s > 0 and show_stratified:
-                from tqdm import tqdm
+            if is_elem := (
+                (is_ad_ref := isinstance(ref_acc, AdRef))
+                or isinstance(ref_acc, LayerAcc | MultiAcc | GraphAcc)
+            ) or (ref_acc is None and X is not self.uns):
+                key = ref_acc.acc if is_ad_ref else ref_acc
+                if isinstance(X, h5py.Dataset) and with_disk:
+                    accumulate[key] += int(np.array(X.shape).prod() * X.dtype.itemsize)
+                elif isinstance(X, BaseCompressedSparseDataset) and with_disk:
+                    accumulate[key] += cs_to_bytes(X._to_backed())
+                elif issparse(X):
+                    accumulate[key] += cs_to_bytes(X)
+                else:
+                    accumulate[key] += X.__sizeof__()
+            if not is_elem or ref_acc is A.X:
+                s = accumulate[ref_acc]
+                if s > 0 and show_stratified:
+                    from tqdm import tqdm
+
+                    print(
+                        f"Size of {repr(ref_acc).replace('A.', '') if ref_acc is not None else 'uns'}: {tqdm.format_sizeof(s, 'B')}"
+                    )
+            return accumulate
 
-                print(
-                    f"Size of {attr.replace('_', '.'):<7}: {tqdm.format_sizeof(s, 'B')}"
-                )
-            sizes[attr] = s
-        return sum(sizes.values())
+        return sum(self.fold(get_size, init=defaultdict(int)).values())
 
     def _gen_repr(self, n_obs, n_vars) -> str:
         backed_at = f" backed at {str(self.filename)!r}" if self.isbacked else ""
@@ -1446,6 +1453,75 @@ def copy(self, filename: PathLike[str] | str | None = None) -> AnnData:
             write_h5ad(filename, self)
             return read_h5ad(filename, backed=mode)
 
+    def fold[T](self, func: FoldFunc[T], *, init: T) -> T:
+        """Accumulate a value starting from init by iterating over the "elems"/leaf nodes of the AnnData object in DFS order.
+
+        Parameters
+        ----------
+        func
+            The function that performs the accumulation
+        init
+            The starting value
+
+
+        Returns
+        -------
+            An accumulated value
+        """
+        accumulate = init
+        for attr_name in [
+            "X",
+            "obs",
+            "var",
+            "obsm",
+            "varm",
+            "obsp",
+            "varp",
+            "layers",
+        ]:
+            attr = getattr(self, attr_name)
+            acc = getattr(A, attr_name)
+            if attr_name != "X":
+                for elem_name in attr:
+                    ref = acc[elem_name]
+                    accumulate = func(
+                        attr[elem_name], accumulate=accumulate, ref_acc=ref
+                    )
+            accumulate = func(attr, accumulate=accumulate, ref_acc=acc)
+        for elem in self.uns:
+            accumulate = func(elem, accumulate=accumulate, ref_acc=None)
+        return accumulate
+
+    def can_write(self, *, store_type: Literal["h5", "zarr"] | None) -> bool:
+        """Whether or not an `AnnData` object can be written to disk for a given store type.
+
+        Parameters
+        ----------
+        store_type
+            Which backing store - `None` indicates that it can be writeable to either.
+
+        Returns
+        -------
+            Whether or not this object is writable.
+        """
+        from anndata._io.specs.registry import _REGISTRY
+
+        writeable_elems = {
+            src_type
+            for (dest_type, src_type, __) in _REGISTRY.write
+            if store_type is None or store_type in dest_type.__module__
+        }
+
+        def predicate(x: RWAble, *, accumulate: bool, ref_acc: AdRef | RefAcc | None):
+            if isinstance(ref_acc, AdRef) or ref_acc is None:
+                if isinstance(x, pd.Series):
+                    # matches behavior in methods.py
+                    x = x._values
+                return accumulate and type(x) in writeable_elems
+            return accumulate
+
+        return self.fold(predicate, init=True)
+
     @deprecated(
         deprecation_msg(
             *("AnnData.concatenate", "anndata.concat"),

src/anndata/types.py

@@ -8,6 +8,9 @@
 
     from array_api.latest import ArrayNamespace
 
+    from anndata.acc import AdRef, RefAcc
+    from anndata.typing import RWAble
+
     from ._core.anndata import AnnData
 
 
@@ -48,3 +51,9 @@ def __dlpack__(
         copy: bool | None = None,
     ) -> Any: ...
     def __dlpack_device__(self) -> tuple[int, int]: ...
+
+
+class FoldFunc[T](Protocol):
+    def __call__(
+        self, elem: RWAble, *, accumulate: T, ref_acc: RefAcc | AdRef | None
+    ) -> T | None: ...

tests/test_readwrite.py

@@ -99,7 +99,7 @@ def dataset_kwargs(request):
 
 
 @pytest.fixture
-def rw(backing_h5ad):
+def rw(backing_h5ad) -> tuple[ad.AnnData, ad.AnnData]:
     M, N = 100, 101
     orig = gen_adata((M, N), **GEN_ADATA_NO_XARRAY_ARGS)
     orig.write(backing_h5ad)
@@ -126,6 +126,27 @@ def dtype(request):
 # ------------------------------------------------------------------------------
 
 
+@pytest.mark.parametrize("store_type", ["h5", "zarr", None])
+def test_can_write(
+    rw: tuple[ad.AnnData, ad.AnnData], store_type: Literal["h5", "zarr"] | None
+):
+    adata, _ = rw
+    assert adata.can_write(store_type=store_type)
+
+
+@pytest.mark.parametrize("store_type", ["h5", "zarr", None])
+def test_can_not_write_with_custom_array(
+    rw: tuple[ad.AnnData, ad.AnnData], store_type: Literal["h5", "zarr"] | None
+):
+    import pyarrow as pa
+
+    adata, _ = rw
+    adata.obs["arrow_array"] = pd.arrays.ArrowExtensionArray(
+        pa.array([{"x": 1, "y": True}] * adata.shape[0])
+    )
+    assert not adata.can_write(store_type=store_type)
+
+
 @pytest.mark.parametrize("typ", ARRAY_TYPES)
 def test_readwrite_roundtrip(typ, tmp_path, diskfmt, diskfmt2):
     pth1 = tmp_path / f"first.{diskfmt}"