Merge pull request #73 from YosefLab/anndata-0.12.0

anndata>=0.12.0 compatibility

Author: colganwi

.gitignore

@@ -2,6 +2,7 @@
 .DS_Store
 *~
 buck-out/
+*._*
 
 # Compiled files
 .venv/

CHANGELOG.md

@@ -8,14 +8,27 @@ and this project adheres to [Semantic Versioning][].
 [keep a changelog]: https://keepachangelog.com/en/1.0.0/
 [semantic versioning]: https://semver.org/spec/v2.0.0.html
 
-## [Unrealeased]
+## [Unreleased]
 
 ### Added
 
 ### Changed
 
 ### Fixed
 
+## [0.2.4] - 2025-11-05
+
+### Added
+
+### Changed
+
+- Deprecated `TreeData.obst_keys` and `TreeData.vart_keys` for consistency with AnnData (https://github.com/scverse/anndata/pull/2093) (#73)
+
+### Fixed
+
+- Added support for zarr v3 (#73)
+- Eliminated deprecations warnings from AnnData>=0.12.0 (#73)
+
 ## [0.2.3] - 2025-10-14
 
 ### Added

pyproject.toml

@@ -4,7 +4,7 @@ requires = [ "hatchling" ]
 
 [project]
 name = "treedata"
-version = "0.2.3"
+version = "0.2.4"
 description = "anndata with trees"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -33,7 +33,7 @@ dependencies = [
   "pathlib",
   "pyarrow",
   "session-info2",
-  "zarr>=2,<3",
+  "zarr>=2",
 ]
 optional-dependencies.dev = [
   "pre-commit",

src/treedata/_core/read.py

@@ -144,6 +144,22 @@ def read_h5ad(
     return read_h5td(filename, backed=backed)
 
 
+def _open_zarr_group(store, *, mode: str = "r") -> tuple[zarr.Group, bool]:
+    """Open a Zarr group and signal whether it should be closed."""
+    if isinstance(store, zarr.Group):
+        return store, False
+    return zarr.open(store, mode=mode), True
+
+
+def _close_zarr_group(group: zarr.Group) -> None:
+    """Close the underlying store for an opened Zarr group if needed."""
+    store = getattr(group, "store", None)
+    if store is not None:
+        close = getattr(store, "close", None)
+        if callable(close):
+            close()
+
+
 def read_zarr(store: str | PathLike | MutableMapping | zarr.Group) -> TreeData:
     """Read from a hierarchical Zarr array store.
 
@@ -152,6 +168,8 @@ def read_zarr(store: str | PathLike | MutableMapping | zarr.Group) -> TreeData:
     store
         The filename, a :class:`~typing.MutableMapping`, or a Zarr storage class.
     """
-    with zarr.open(store, mode="r") as f:
-        d = _read_tdata(f, store, backed=False)
+    group, should_close = _open_zarr_group(store)
+    d = _read_tdata(group, store, backed=False)
+    if should_close:
+        _close_zarr_group(group)
     return TreeData(**d)

src/treedata/_core/treedata.py

@@ -232,10 +232,21 @@ def _init_as_view(self, tdata_ref: TreeData, oidx: Index1D | None, vidx: Index1D
 
     def obst_keys(self) -> list[str]:
         """List keys of variable annotation `obst`."""
+        # add future warning
+        warnings.warn(
+            "Use obst (e.g. `k in tdata.obst` or `adata.obst.keys()`) instead of TreeData.obst_keys, TreeData.obst_keys is deprecated and will be removed in the future.",
+            FutureWarning,
+            stacklevel=2,
+        )
         return list(self._obst.keys())
 
     def vart_keys(self) -> list[str]:
         """List keys of variable annotation `vart`."""
+        warnings.warn(
+            "Use vart (e.g. `k in tdata.vart` or `adata.vart.keys()`) instead of TreeData.vart_keys, TreeData.vart_keys is deprecated and will be removed in the future.",
+            FutureWarning,
+            stacklevel=2,
+        )
         return list(self._vart.keys())
 
     @property
@@ -496,7 +507,7 @@ def write_h5td(
     def write_zarr(
         self,
         store: MutableMapping | PathLike,
-        chunks: bool | int | tuple[int, ...] | None = None,
+        chunks: tuple[int, ...] | None = None,
         **kwargs,
     ):
         """Write a hierarchical Zarr array store.
@@ -510,6 +521,13 @@ def write_zarr(
         """
         from .write import write_zarr
 
+        if isinstance(chunks, bool):
+            msg = (
+                "Passing `write_zarr(adata, chunks=True)` is no longer supported. "
+                "Please pass `write_zarr(adata)` instead."
+            )
+            raise ValueError(msg)
+
         write_zarr(store, self, chunks=chunks)
 
     def to_memory(self, copy=False) -> TreeData:
@@ -520,7 +538,7 @@ def to_memory(self, copy=False) -> TreeData:
             copy:
                 Whether the arrays that are already in-memory should be copied.
         """
-        adata = super().to_memory(copy)
+        adata = super().to_memory(copy=copy)
         tdata = TreeData(
             adata,
             obst=self.obst.copy(),

src/treedata/_core/write.py

@@ -3,6 +3,7 @@
 import json
 import warnings
 from collections.abc import MutableMapping
+from importlib.metadata import version as get_version
 from os import PathLike
 from pathlib import Path
 from typing import Any, Literal
@@ -18,8 +19,9 @@
 from treedata._core.aligned_mapping import AxisTrees
 from treedata._core.treedata import TreeData
 
-ANDATA_VERSION = version.parse(ad.__version__)
-USE_EXPERIMENTAL = ANDATA_VERSION < version.parse("0.11.0")
+ANNDATA_VERSION = version.parse(get_version("anndata"))
+USE_EXPERIMENTAL = ANNDATA_VERSION < version.parse("0.11.0")
+ZARR_V2 = version.parse(get_version("zarr")) < version.parse("3.0.0")
 
 
 def _make_serializable(data: Any) -> Any:
@@ -63,7 +65,7 @@ def _serialize_axis_trees(trees: AxisTrees) -> str:
     return json.dumps(_make_serializable(d))
 
 
-def _write_tdata(f, tdata, filename, **kwargs) -> None:
+def _write_tdata(f, tdata, filename, chunks=None, **kwargs) -> None:
     """Write TreeData to file."""
     # Add encoding type and version
     f = f["/"]
@@ -73,7 +75,7 @@ def _write_tdata(f, tdata, filename, **kwargs) -> None:
     tdata.strings_to_categoricals()
     # Write X if not backed
     if not (tdata.isbacked and Path(tdata.filename) == Path(filename)):
-        _write_elem(f, "X", tdata.X, dataset_kwargs=kwargs)
+        _write_elem(f, "X", tdata.X, dataset_kwargs=kwargs.update({"chunks": chunks}) if chunks else kwargs)
     # Write array elements
     for key in ["obs", "var", "label", "allow_overlap", "alignment"]:
         _write_elem(f, key, getattr(tdata, key), dataset_kwargs=kwargs)
@@ -152,7 +154,32 @@ def write_h5ad(
     )
 
 
-def write_zarr(filename: MutableMapping | PathLike, tdata: TreeData, **kwargs) -> None:
+def _open_zarr_group(
+    store: MutableMapping | PathLike | zarr.Group,
+    *,
+    mode: str = "w",
+) -> tuple[zarr.Group, bool]:
+    """Open a Zarr group for writing and indicate ownership."""
+    if isinstance(store, zarr.Group):
+        return store, False
+    group_kwargs: dict[str, int] = {}
+    if not ZARR_V2:
+        group_kwargs["zarr_format"] = 3
+    return zarr.open_group(store, mode=mode, **group_kwargs), True
+
+
+def _close_zarr_group(group: zarr.Group) -> None:
+    """Close the underlying store for an opened Zarr group if needed."""
+    store = getattr(group, "store", None)
+    if store is not None:
+        close = getattr(store, "close", None)
+        if callable(close):
+            close()
+
+
+def write_zarr(
+    filename: MutableMapping | PathLike | zarr.Group, tdata: TreeData, chunks: tuple[int, ...] | None = None, **kwargs
+) -> None:
     """Write `.zarr`-formatted zarr file.
 
     Parameters
@@ -164,5 +191,11 @@ def write_zarr(filename: MutableMapping | PathLike, tdata: TreeData, **kwargs) -
     kwargs
         Additional keyword arguments passed to :func:`zarr.save`.
     """
-    with zarr.open(filename, mode="w") as f:
-        _write_tdata(f, tdata, filename, **kwargs)
+    # if not ZARR_V2 and "zarr_format" not in kwargs:
+    #    f = zarr.open(store=filename, mode="w", zarr_format=3)
+    # else:
+    print(kwargs)
+    f, should_close = _open_zarr_group(filename)
+    _write_tdata(f, tdata, filename, chunks, **kwargs)
+    if should_close:
+        _close_zarr_group(f)

tests/test_base.py

@@ -246,7 +246,7 @@ def test_transpose(adata, tree):
     treedata_transpose = treedata.transpose()
     assert np.array_equal(treedata.X.T, treedata_transpose.X)
     assert treedata.obst["tree"].nodes == treedata_transpose.vart["tree"].nodes
-    assert treedata_transpose.obst_keys() == []
+    assert list(treedata_transpose.obst.keys()) == []
     assert np.array_equal(treedata.obs_names, treedata.T.obs_names)
 
 

tests/test_readwrite.py

@@ -102,7 +102,7 @@ def test_read_anndata(X, tmp_path):
     assert tdata.label == "tree"
     assert tdata.allow_overlap is True
     assert tdata.alignment == "leaves"
-    assert tdata.obst_keys() == []
+    assert list(tdata.obst.keys()) == []
 
 
 def test_read_no_X(X, tmp_path):