Update zarr io (#66)

* Update `write_zarr` argument names and type hints to accommodate
MuData and TreeData.

* Refactor `read_zarr` to support multiple backends and update
docstrings

Author: WeilerP

src/fancypackage/io/_zarr.py

@@ -1,19 +1,26 @@
+from __future__ import annotations
+
 import warnings
 from pathlib import Path
+from typing import Literal, TYPE_CHECKING
 
 import zarr
 
 import anndata as ad
 from anndata import AnnData
 
+if TYPE_CHECKING:
+    from mudata import MuData
+    from treedata import TreeData
+
 
-def write_zarr(adata: AnnData, path: Path) -> None:
-    """Write AnnData to hierarchical Zarr array zip store.
+def write_zarr(data: AnnData | MuData | TreeData, path: Path) -> None:
+    """Write AnnData-like object to hierarchical Zarr array zip store.
 
     Parameters
     ----------
     adata
-        AnnData to save.
+        AnnData-like object to save; supports AnnData, MuData and TreeData.
     path
         Filename of Zarr store.
 
@@ -24,24 +31,36 @@ def write_zarr(adata: AnnData, path: Path) -> None:
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", UserWarning)
         store = zarr.storage.ZipStore(path, mode="w")
-        adata.write_zarr(store=store)
+        data.write_zarr(store=store)
         store.close()
 
 
-def read_zarr(path: Path) -> AnnData:
-    """Read hierarchical Zarr array zip store to AnnData.
+def read_zarr(path: Path, backend: Literal["anndata", "mudata", "treedata"] = "anndata") -> AnnData | MuData | TreeData:
+    """Read hierarchical Zarr array zip store to AnnData-like object.
 
     Parameters
     ----------
     path
         Filename of Zarr store.
+    backend
+        Backend to use for reading the Zarr store. Supported backends are "anndata" (default), "mudata" and "treedata".
 
     Returns
     -------
-    Read AnnData object.
+    The read AnnData-like object (Anndata, MuData, or TreeData).
     """
     store = zarr.storage.ZipStore(path, mode="r")
-    adata = ad.io.read_zarr(store)
+
+    if backend == "anndata":
+        data = ad.read_zarr(store)
+    elif backend == "mudata":
+        import mudata as md
+
+        data = md.read_zarr(store)
+    elif backend == "treedata":
+        import treedata as td
+
+        data = td.read_zarr(store)
     store.close()
 
-    return adata
+    return data