From 333a9e3220841debbcd4472d43f0347e64389f2e Mon Sep 17 00:00:00 2001
From: Joe Hamman <joe@earthmover.io>
Date: Sat, 7 Dec 2024 09:31:54 -0800
Subject: [PATCH] [v3] implement / deprecate zarr.tree (#2537)

---
 src/zarr/api/asynchronous.py | 26 ++++++++++++++++++++++++--
 src/zarr/api/synchronous.py  |  7 +++++--
 src/zarr/core/group.py       |  2 +-
 tests/test_api.py            |  7 ++++---
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py
index e5d09f8c3a..26822f725b 100644
--- a/src/zarr/api/asynchronous.py
+++ b/src/zarr/api/asynchronous.py
@@ -7,6 +7,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from typing_extensions import deprecated
 
 from zarr.core.array import Array, AsyncArray, get_array_metadata
 from zarr.core.buffer import NDArrayLike
@@ -493,8 +494,29 @@ async def save_group(
     await asyncio.gather(*aws)
 
 
-async def tree(*args: Any, **kwargs: Any) -> None:
-    raise NotImplementedError
+@deprecated("Use AsyncGroup.tree instead.")
+async def tree(grp: AsyncGroup, expand: bool | None = None, level: int | None = None) -> Any:
+    """Provide a rich display of the hierarchy.
+
+    Parameters
+    ----------
+    grp : Group
+        Zarr or h5py group.
+    expand : bool, optional
+        Only relevant for HTML representation. If True, tree will be fully expanded.
+    level : int, optional
+        Maximum depth to descend into hierarchy.
+
+    Returns
+    -------
+    TreeRepr
+        A pretty-printable object displaying the hierarchy.
+
+    .. deprecated:: 3.0.0
+        `zarr.tree()` is deprecated and will be removed in a future release.
+        Use `group.tree()` instead.
+    """
+    return await grp.tree(expand=expand, level=level)
 
 
 async def array(
diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py
index 9616c41355..8e8ecf40b8 100644
--- a/src/zarr/api/synchronous.py
+++ b/src/zarr/api/synchronous.py
@@ -2,6 +2,8 @@
 
 from typing import TYPE_CHECKING, Any, Literal
 
+from typing_extensions import deprecated
+
 import zarr.api.asynchronous as async_api
 from zarr._compat import _deprecate_positional_args
 from zarr.core.array import Array, AsyncArray
@@ -155,8 +157,9 @@ def save_group(
     )
 
 
-def tree(*args: Any, **kwargs: Any) -> None:
-    return sync(async_api.tree(*args, **kwargs))
+@deprecated("Use Group.tree instead.")
+def tree(grp: Group, expand: bool | None = None, level: int | None = None) -> Any:
+    return sync(async_api.tree(grp._async_group, expand=expand, level=level))
 
 
 # TODO: add type annotations for kwargs
diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py
index 2ca6e209fd..adc6399276 100644
--- a/src/zarr/core/group.py
+++ b/src/zarr/core/group.py
@@ -1449,7 +1449,7 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An
         from zarr.core._tree import group_tree_async
 
         if expand is not None:
-            raise NotImplementedError("'expanded' is not yet implemented.")
+            raise NotImplementedError("'expand' is not yet implemented.")
         return await group_tree_async(self, max_depth=level)
 
     async def empty(
diff --git a/tests/test_api.py b/tests/test_api.py
index 11977e8e32..90f6dae110 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -294,15 +294,16 @@ def test_load_array(memory_store: Store) -> None:
 
 
 def test_tree() -> None:
+    pytest.importorskip("rich")
     g1 = zarr.group()
     g1.create_group("foo")
     g3 = g1.create_group("bar")
     g3.create_group("baz")
     g5 = g3.create_group("qux")
     g5.create_array("baz", shape=100, chunks=10)
-    # TODO: complete after tree has been reimplemented
-    # assert repr(zarr.tree(g1)) == repr(g1.tree())
-    # assert str(zarr.tree(g1)) == str(g1.tree())
+    with pytest.warns(DeprecationWarning):
+        assert repr(zarr.tree(g1)) == repr(g1.tree())
+        assert str(zarr.tree(g1)) == str(g1.tree())
 
 
 # @pytest.mark.parametrize("stores_from_path", [False, True])