fix: added support and tests for opening TreeStore

Author: bossbeagle1509

src/blosc2/schunk.py

@@ -1471,12 +1471,18 @@ def __dealloc__(self):
 
 
 def _open_special_store(urlpath, mode, offset, **kwargs):
-    if urlpath.endswith(".b2z") or urlpath.endswith(".b2d"):
+    if urlpath.endswith(".b2d"):
         if offset != 0:
             raise ValueError("Offset must be 0 for DictStore")
         from blosc2.dict_store import DictStore
 
         return DictStore(urlpath, mode=mode, **kwargs)
+    elif urlpath.endswith(".b2z"):
+        if offset != 0:
+            raise ValueError("Offset must be 0 for TreeStore")
+        from blosc2.tree_store import TreeStore
+
+        return TreeStore(urlpath, mode=mode, **kwargs)
     elif urlpath.endswith(".b2e"):
         if offset != 0:
             raise ValueError("Offset must be 0 for EmbedStore")
@@ -1520,9 +1526,18 @@ def _process_opened_object(res):
 
 def open(
     urlpath: str | pathlib.Path | blosc2.URLPath, mode: str = "a", offset: int = 0, **kwargs: dict
-) -> blosc2.SChunk | blosc2.NDArray | blosc2.C2Array | blosc2.LazyArray | blosc2.Proxy | Any:
+) -> (
+    blosc2.SChunk
+    | blosc2.NDArray
+    | blosc2.C2Array
+    | blosc2.LazyArray
+    | blosc2.Proxy
+    | blosc2.DictStore
+    | blosc2.TreeStore
+    | blosc2.EmbedStore
+):
     """Open a persistent :ref:`SChunk`, :ref:`NDArray`, a remote :ref:`C2Array`,
-    a :ref:`Proxy`, a :ref:`DictStore` or an :ref:`EmbedStore`.
+    a :ref:`Proxy`, a :ref:`DictStore`, :ref:`EmbedStore`, or :ref:`TreeStore`.
 
     See the `Notes` section for more info on opening `Proxy` objects.
 
@@ -1558,7 +1573,7 @@ def open(
 
     Returns
     -------
-    out: :ref:`SChunk`, :ref:`NDArray`, :ref:`C2Array`, :ref:`DictStore` or :ref:`EmbedStore`
+    out: :ref:`SChunk`, :ref:`NDArray`, :ref:`C2Array`, :ref:`DictStore`, :ref:`EmbedStore`, or :ref:`TreeStore`
         The object found in the path.
 
     Notes

tests/test_tree_store.py

@@ -919,3 +919,18 @@ def test_key_normalization():
         assert "/group/data2" in tstore
 
     os.remove("test_key_normalization.b2z")
+
+
+def test_open_context_manager(populated_tree_store):
+    """Test opening via blosc2.open as a context manager."""
+    tstore_fixture, path = populated_tree_store
+    if ".b2d" in path:
+        pytest.skip("This test is only for b2z storage")
+    # Close the fixture store to ensure data is written to disk
+    tstore_fixture.close()
+
+    # Test opening via blosc2.open as a context manager
+    with blosc2.open(path, mode="r") as tstore:
+        assert isinstance(tstore, TreeStore)
+        assert "/child0/data" in tstore
+        assert np.array_equal(tstore["/child0/data"][:], np.array([1, 2, 3]))