Use metadata-based DictStore discovery and warn on leaf mismatches

Author: FrancescAlted

src/blosc2/dict_store.py

@@ -10,6 +10,7 @@
 import os
 import shutil
 import tempfile
+import warnings
 import zipfile
 from typing import TYPE_CHECKING, Any
 
@@ -94,6 +95,9 @@ class DictStore:
     -----
     - External persistence uses the following file extensions:
       .b2nd for NDArray, .b2f for SChunk, and .b2b for BatchStore.
+      These suffixes are a naming convention for newly written leaves; when
+      reopening an existing store, leaf typing is resolved from object
+      metadata instead of trusting the suffix alone.
     """
 
     def __init__(
@@ -112,7 +116,7 @@ def __init__(
         """
         See :class:`DictStore` for full documentation of parameters.
         """
-        self.localpath = localpath if isinstance(localpath, (str, bytes)) else str(localpath)
+        self.localpath = localpath if isinstance(localpath, str | bytes) else str(localpath)
         if not self.localpath.endswith((".b2z", ".b2d")):
             raise ValueError(f"localpath must have a .b2z or .b2d extension; you passed: {self.localpath}")
         if mode not in ("r", "w", "a"):
@@ -182,13 +186,7 @@ def _init_read_mode(self, dparams: blosc2.DParams | None = None):
                 mmap_mode=self.mmap_mode,
                 dparams=dparams,
             )
-            for filepath in self.offsets:
-                if filepath.endswith((".b2nd", ".b2f", ".b2b")):
-                    if filepath.endswith(".b2nd"):
-                        key = "/" + filepath[:-5]
-                    else:
-                        key = "/" + filepath[:-4]
-                    self.map_tree[key] = filepath
+            self._update_map_tree_from_offsets()
         else:  # .b2d
             if not os.path.isdir(self.localpath):
                 raise FileNotFoundError(f"Directory {self.localpath} does not exist for reading.")
@@ -204,6 +202,90 @@ def _init_read_mode(self, dparams: blosc2.DParams | None = None):
         self._estore = EmbedStore(_from_schunk=schunk)
         self.storage.meta = self._estore.storage.meta
 
+    @staticmethod
+    def _logical_key_from_relpath(rel_path: str) -> str:
+        """Map an external leaf path to its logical tree key."""
+        rel_path = rel_path.replace(os.sep, "/")
+        key = os.path.splitext(rel_path)[0]
+        if not key.startswith("/"):
+            key = "/" + key
+        return key
+
+    @staticmethod
+    def _expected_ext_from_kind(kind: str) -> str:
+        """Return the canonical write-time suffix for a supported external leaf kind."""
+        if kind == "ndarray":
+            return ".b2nd"
+        if kind == "batchstore":
+            return ".b2b"
+        return ".b2f"
+
+    @classmethod
+    def _opened_external_kind(
+        cls,
+        opened: blosc2.NDArray | SChunk | blosc2.VLArray | blosc2.BatchStore | C2Array | Any,
+        rel_path: str,
+    ) -> str | None:
+        """Return the supported external leaf kind for an already opened object."""
+        processed = _process_opened_object(opened)
+        if isinstance(processed, blosc2.BatchStore):
+            kind = "batchstore"
+        elif isinstance(processed, blosc2.VLArray):
+            kind = "vlarray"
+        elif isinstance(processed, blosc2.NDArray):
+            kind = "ndarray"
+        elif isinstance(processed, SChunk):
+            kind = "schunk"
+        else:
+            warnings.warn(
+                f"Ignoring unsupported Blosc2 object at '{rel_path}' during DictStore discovery: "
+                f"{type(processed).__name__}",
+                UserWarning,
+                stacklevel=2,
+            )
+            return None
+
+        expected_ext = cls._expected_ext_from_kind(kind)
+        found_ext = os.path.splitext(rel_path)[1]
+        if found_ext != expected_ext:
+            warnings.warn(
+                f"External leaf '{rel_path}' uses extension '{found_ext}' but metadata resolves to "
+                f"{type(processed).__name__}; expected '{expected_ext}'.",
+                UserWarning,
+                stacklevel=2,
+            )
+        return kind
+
+    def _probe_external_leaf_path(self, rel_path: str) -> bool:
+        """Return whether a working-dir file is a supported external leaf."""
+        urlpath = os.path.join(self.working_dir, rel_path)
+        try:
+            opened = blosc2.blosc2_ext.open(
+                urlpath,
+                mode="r",
+                offset=0,
+                mmap_mode=self.mmap_mode,
+                dparams=self.dparams,
+            )
+        except Exception:
+            return False
+        return self._opened_external_kind(opened, rel_path) is not None
+
+    def _probe_external_leaf_offset(self, filepath: str) -> bool:
+        """Return whether a zip member is a supported external leaf."""
+        offset = self.offsets[filepath]["offset"]
+        try:
+            opened = blosc2.blosc2_ext.open(
+                self.b2z_path,
+                mode="r",
+                offset=offset,
+                mmap_mode=self.mmap_mode,
+                dparams=self.dparams,
+            )
+        except Exception:
+            return False
+        return self._opened_external_kind(opened, filepath) is not None
+
     def _init_write_append_mode(
         self,
         cparams: blosc2.CParams | None,
@@ -229,24 +311,23 @@ def _init_write_append_mode(
         self._update_map_tree()
 
     def _update_map_tree(self):
-        # Build map_tree from .b2nd and .b2f files in working dir
+        # Build map_tree from supported external leaves in working dir.
         for root, _, files in os.walk(self.working_dir):
             for file in files:
                 filepath = os.path.join(root, file)
-                if filepath.endswith((".b2nd", ".b2f", ".b2b")):
-                    # Convert filename to key: remove extension and ensure starts with /
-                    rel_path = os.path.relpath(filepath, self.working_dir)
-                    # Normalize path separators to forward slashes for cross-platform consistency
-                    rel_path = rel_path.replace(os.sep, "/")
-                    if rel_path.endswith(".b2nd"):
-                        key = rel_path[:-5]
-                    elif rel_path.endswith(".b2b") or rel_path.endswith(".b2f"):
-                        key = rel_path[:-4]
-                    else:
-                        continue
-                    if not key.startswith("/"):
-                        key = "/" + key
-                    self.map_tree[key] = rel_path
+                if os.path.abspath(filepath) == os.path.abspath(self.estore_path):
+                    continue
+                rel_path = os.path.relpath(filepath, self.working_dir).replace(os.sep, "/")
+                if self._probe_external_leaf_path(rel_path):
+                    self.map_tree[self._logical_key_from_relpath(rel_path)] = rel_path
+
+    def _update_map_tree_from_offsets(self):
+        """Build map_tree from supported external leaves in a zip store."""
+        for filepath in self.offsets:
+            if filepath == "embed.b2e":
+                continue
+            if self._probe_external_leaf_offset(filepath):
+                self.map_tree[self._logical_key_from_relpath(filepath)] = filepath
 
     @property
     def estore(self) -> EmbedStore:
@@ -255,13 +336,13 @@ def estore(self) -> EmbedStore:
 
     @staticmethod
     def _value_nbytes(value: blosc2.Array | SChunk | blosc2.VLArray | blosc2.BatchStore) -> int:
-        if isinstance(value, (blosc2.VLArray, blosc2.BatchStore)):
+        if isinstance(value, blosc2.VLArray | blosc2.BatchStore):
             return value.schunk.nbytes
         return value.nbytes
 
     @staticmethod
     def _is_external_value(value: blosc2.Array | SChunk | blosc2.VLArray | blosc2.BatchStore) -> bool:
-        return isinstance(value, (blosc2.NDArray, SChunk, blosc2.VLArray, blosc2.BatchStore)) and bool(
+        return isinstance(value, blosc2.NDArray | SChunk | blosc2.VLArray | blosc2.BatchStore) and bool(
             getattr(value, "urlpath", None)
         )
 
@@ -406,12 +487,14 @@ def values(self) -> Iterator[blosc2.NDArray | SChunk | C2Array]:
                 if self.is_zip_store:
                     if filepath in self.offsets:
                         offset = self.offsets[filepath]["offset"]
-                        yield blosc2.blosc2_ext.open(
-                            self.b2z_path,
-                            mode="r",
-                            offset=offset,
-                            mmap_mode=self.mmap_mode,
-                            dparams=self.dparams,
+                        yield _process_opened_object(
+                            blosc2.blosc2_ext.open(
+                                self.b2z_path,
+                                mode="r",
+                                offset=offset,
+                                mmap_mode=self.mmap_mode,
+                                dparams=self.dparams,
+                            )
                         )
                 else:
                     urlpath = os.path.join(self.working_dir, filepath)
@@ -438,12 +521,14 @@ def items(self) -> Iterator[tuple[str, blosc2.NDArray | SChunk | C2Array]]:
                         offset = self.offsets[filepath]["offset"]
                         yield (
                             key,
-                            blosc2.blosc2_ext.open(
-                                self.b2z_path,
-                                mode="r",
-                                offset=offset,
-                                mmap_mode=self.mmap_mode,
-                                dparams=self.dparams,
+                            _process_opened_object(
+                                blosc2.blosc2_ext.open(
+                                    self.b2z_path,
+                                    mode="r",
+                                    offset=offset,
+                                    mmap_mode=self.mmap_mode,
+                                    dparams=self.dparams,
+                                )
                             ),
                         )
                 else:

tests/test_dict_store.py

@@ -16,6 +16,22 @@
 from blosc2.dict_store import DictStore
 
 
+def _rename_store_member(store_path, old_name, new_name):
+    """Rename an external leaf inside a .b2d/.b2z store without changing its contents."""
+    if str(store_path).endswith(".b2d"):
+        old_path = os.path.join(store_path, old_name.replace("/", os.sep))
+        new_path = os.path.join(store_path, new_name.replace("/", os.sep))
+        os.rename(old_path, new_path)
+        return
+
+    tmp_zip = f"{store_path}.tmp"
+    with zipfile.ZipFile(store_path, "r") as src, zipfile.ZipFile(tmp_zip, "w", zipfile.ZIP_STORED) as dst:
+        for info in src.infolist():
+            arcname = new_name if info.filename == old_name else info.filename
+            dst.writestr(arcname, src.read(info.filename), compress_type=zipfile.ZIP_STORED)
+    os.replace(tmp_zip, store_path)
+
+
 @pytest.fixture(params=["b2d", "b2z"])
 def populated_dict_store(request):
     """Create and populate a DictStore for tests.
@@ -266,6 +282,78 @@ def test_external_vlarray_file_and_reopen(tmp_path):
         assert value.vlmeta["description"] == "External VLArray"
 
 
+@pytest.mark.parametrize("storage_type", ["b2d", "b2z"])
+def test_metadata_discovery_reopens_renamed_external_ndarray(storage_type, tmp_path):
+    path = tmp_path / f"test_renamed_ndarray.{storage_type}"
+    ext_path = tmp_path / "renamed_array_source.b2nd"
+
+    with DictStore(str(path), mode="w", threshold=None) as dstore:
+        arr_external = blosc2.arange(5, urlpath=str(ext_path), mode="w")
+        arr_external.vlmeta["description"] = "Renamed NDArray"
+        dstore["/dir1/node3"] = arr_external
+
+    old_name = "dir1/node3.b2nd"
+    new_name = "dir1/node3.weird"
+    _rename_store_member(str(path), old_name, new_name)
+
+    with pytest.warns(UserWarning, match=r"node3\.weird'.*NDArray.*expected '\.b2nd'"):
+        dstore_read = DictStore(str(path), mode="r")
+    with dstore_read:
+        assert dstore_read.map_tree["/dir1/node3"] == new_name
+        node3 = dstore_read["/dir1/node3"]
+        assert isinstance(node3, blosc2.NDArray)
+        assert np.array_equal(node3[:], np.arange(5))
+        assert node3.vlmeta["description"] == "Renamed NDArray"
+
+
+@pytest.mark.parametrize("storage_type", ["b2d", "b2z"])
+def test_metadata_discovery_reopens_renamed_external_vlarray(storage_type, tmp_path):
+    path = tmp_path / f"test_renamed_vlarray.{storage_type}"
+    ext_path = tmp_path / "renamed_vlarray_source.b2frame"
+    values = ["alpha", {"nested": True}, None, (1, 2, 3)]
+
+    vlarray = blosc2.VLArray(urlpath=str(ext_path), mode="w", contiguous=True)
+    vlarray.extend(values)
+    vlarray.vlmeta["description"] = "Renamed VLArray"
+
+    with DictStore(str(path), mode="w", threshold=None) as dstore:
+        dstore["/dir1/vlarray_ext"] = vlarray
+
+    old_name = "dir1/vlarray_ext.b2f"
+    new_name = "dir1/vlarray_ext.renamed"
+    _rename_store_member(str(path), old_name, new_name)
+
+    with pytest.warns(UserWarning, match=r"vlarray_ext\.renamed'.*VLArray.*expected '\.b2f'"):
+        dstore_read = DictStore(str(path), mode="r")
+    with dstore_read:
+        assert dstore_read.map_tree["/dir1/vlarray_ext"] == new_name
+        value = dstore_read["/dir1/vlarray_ext"]
+        assert isinstance(value, blosc2.VLArray)
+        assert list(value) == values
+        assert value.vlmeta["description"] == "Renamed VLArray"
+
+
+def test_metadata_discovery_warns_and_skips_unsupported_blosc2_leaf(tmp_path):
+    path = tmp_path / "test_unsupported_lazyexpr.b2d"
+
+    with DictStore(str(path), mode="w") as dstore:
+        dstore["/embedded"] = np.arange(3)
+
+    a = blosc2.asarray(np.arange(5), urlpath=str(tmp_path / "a.b2nd"), mode="w")
+    b = blosc2.asarray(np.arange(5), urlpath=str(tmp_path / "b.b2nd"), mode="w")
+    expr = a + b
+    expr_path = path / "unsupported_lazyexpr.b2nd"
+    expr.save(str(expr_path))
+
+    with pytest.warns(
+        UserWarning, match=r"Ignoring unsupported Blosc2 object.*unsupported_lazyexpr\.b2nd.*LazyExpr"
+    ):
+        dstore_read = DictStore(str(path), mode="r")
+    with dstore_read:
+        assert "/unsupported_lazyexpr" not in dstore_read
+        assert "/embedded" in dstore_read
+
+
 def _digest_value(value):
     """Return a bytes digest of a stored value."""
     if isinstance(value, blosc2.SChunk):

tests/test_tree_store.py

@@ -7,6 +7,7 @@
 
 import os
 import shutil
+import zipfile
 
 import numpy as np
 import pytest
@@ -15,6 +16,22 @@
 from blosc2.tree_store import TreeStore
 
 
+def _rename_store_member(store_path, old_name, new_name):
+    """Rename an external leaf inside a .b2d/.b2z store without changing its contents."""
+    if str(store_path).endswith(".b2d"):
+        old_path = os.path.join(store_path, old_name.replace("/", os.sep))
+        new_path = os.path.join(store_path, new_name.replace("/", os.sep))
+        os.rename(old_path, new_path)
+        return
+
+    tmp_zip = f"{store_path}.tmp"
+    with zipfile.ZipFile(store_path, "r") as src, zipfile.ZipFile(tmp_zip, "w", zipfile.ZIP_STORED) as dst:
+        for info in src.infolist():
+            arcname = new_name if info.filename == old_name else info.filename
+            dst.writestr(arcname, src.read(info.filename), compress_type=zipfile.ZIP_STORED)
+    os.replace(tmp_zip, store_path)
+
+
 @pytest.fixture(params=["b2d", "b2z"])
 def populated_tree_store(request):
     """A fixture that creates and populates a TreeStore."""
@@ -671,6 +688,28 @@ def test_external_batchstore_support(tmp_path):
         assert [batch[:] for batch in retrieved] == [[{"id": 1}, {"id": 2}], [{"id": 3}]]
 
 
+@pytest.mark.parametrize("storage_type", ["b2d", "b2z"])
+def test_metadata_discovery_reopens_renamed_batchstore_leaf(storage_type, tmp_path):
+    store_path = tmp_path / f"test_batchstore_renamed.{storage_type}"
+
+    with TreeStore(str(store_path), mode="w", threshold=0) as tstore:
+        bstore = blosc2.BatchStore(max_blocksize=2)
+        bstore.extend([[{"id": 1}, {"id": 2}], [{"id": 3}]])
+        tstore["/data/batchstore"] = bstore
+
+    old_name = "data/batchstore.b2b"
+    new_name = "data/batchstore.odd"
+    _rename_store_member(str(store_path), old_name, new_name)
+
+    with pytest.warns(UserWarning, match=r"batchstore\.odd'.*BatchStore.*expected '\.b2b'"):
+        tstore = TreeStore(str(store_path), mode="r")
+    with tstore:
+        assert tstore.map_tree["/data/batchstore"] == new_name
+        retrieved = tstore["/data/batchstore"]
+        assert isinstance(retrieved, blosc2.BatchStore)
+        assert [batch[:] for batch in retrieved] == [[{"id": 1}, {"id": 2}], [{"id": 3}]]
+
+
 def test_treestore_vlmeta_externalized_b2d(tmp_path):
     store_path = tmp_path / "test_vlmeta_externalized.b2d"